完整部署快照

ixformer_sdk/train/speedformer/policy/gpt2.py

@@ -0,0 +1,27 @@
+import torch
+import torch.nn as nn
+from torch.nn import LayerNorm
+from types import ModuleType, MethodType
+from abc import ABC
+
+from ixformer.train.speedformer.models.gpt2.modeling_gpt2 import GPT2FlashAttention2
+
+from ixformer.train.speedformer.layers.normalization import replace_layernorm_forward
+from ixformer.train.speedformer.layers.gpt2.attention import replace_flash_attn_forward
+
+
+class GPT2Replacer(ABC):
+    def __init__(self) -> None:
+        super().__init__()
+
+    @staticmethod
+    def accelerate(model):
+        # layer/kernel replace
+        for name, module in model.named_modules():
+            if isinstance(module, LayerNorm):
+                module.forward = MethodType(replace_layernorm_forward, module)
+            if isinstance(module, GPT2FlashAttention2):
+                module._flash_attention_forward = MethodType(
+                    replace_flash_attn_forward, module)
+
+        return model