[WIP] Adding support for FP8 training

open_lm/main.py

@@ -63,7 +63,6 @@
     terminate_sync_process,
 )
 
-
 LATEST_CHECKPOINT_NAME = "epoch_latest.pt"
 
 
@@ -439,6 +438,8 @@ def main(args):
 
     random_seed(args.seed, args.rank)
 
+    all_gpus = dist.new_group(backend="nccl")
+
     if args.distributed:
         if args.fsdp:
             transformer_layer_cls = None
@@ -498,12 +499,14 @@ def main(args):
             random_seed(args.seed, rank=0)
             model = FSDP(
                 model,
+                process_group=all_gpus,
                 auto_wrap_policy=transformer_auto_wrapper_policy,
                 device_id=device,
                 mixed_precision=mp_policy,
                 cpu_offload=CPUOffload(offload_params=args.fsdp_cpu_offload),
                 use_orig_params=args.fsdp_use_orig_params,
                 limit_all_gathers=args.fsdp_limit_all_gathers,
+                sync_module_states=True,
                 **fsdp_kwargs,
             )
 
@@ -754,6 +757,7 @@ def main(args):
             total_steps=total_steps,
             args=args,
             tb_writer=writer,
+            all_gpus=all_gpus,
         )
 
         if args.distributed:

open_lm/model.py

@@ -35,6 +35,18 @@
 except ImportError:
     MambaLMHeadModel = None
 
+# Adding flag if using TE FP8
+using_te = False
+try:
+    import transformer_engine.pytorch as te
+    from transformer_engine.common import recipe
+
+    fp8_format = recipe.Format.HYBRID
+    fp8_recipe = recipe.DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max")
+    using_te = True
+except ImportError as ie:
+    using_te = False
+
 # from openclip
 _MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
 _MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
@@ -117,41 +129,72 @@ def __init__(self, layer_id, args: Params):
         super().__init__()
         self.n_heads = args.n_heads
         self.head_dim = args.dim // args.n_heads
-        self.in_proj = nn.Linear(args.dim, 3 * args.n_heads * self.head_dim, bias=False)
-        self.out_proj = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
+        if using_te:
+            self.in_proj = te.Linear(args.dim, 3 * args.n_heads * self.head_dim, bias=False, device="cuda")
+            self.out_proj = te.Linear(args.n_heads * self.head_dim, args.dim, bias=False, device="cuda")
+        else:
+            self.in_proj = nn.Linear(args.dim, 3 * args.n_heads * self.head_dim, bias=False)
+            self.out_proj = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
         self.pos_embed = get_pos_embed(args)
         self.attn_fn = args.attn_func
         self.apply_qk_norm = args.apply_qk_norm
 
-        # initialize norm layers for queries and keys if needed
-        self.q_norm = (
-            args.norm_type(
-                args.n_heads * self.head_dim,
-                eps=args.norm_eps,
+        if using_te:
+            # initialize norm layers for queries and keys if needed
+            self.q_norm = (
+                te.LayerNorm(
+                    args.n_heads * self.head_dim,
+                    eps=args.norm_eps,
+                )
+                if self.apply_qk_norm
+                else nn.Identity()
             )
-            if self.apply_qk_norm
-            else nn.Identity()
-        )
-        self.k_norm = (
-            args.norm_type(
-                args.n_heads * self.head_dim,
-                eps=args.norm_eps,
+            self.k_norm = (
+                te.LayerNorm(
+                    args.n_heads * self.head_dim,
+                    eps=args.norm_eps,
+                )
+                if self.apply_qk_norm
+                else nn.Identity()
+            )
+        else:
+            # initialize norm layers for queries and keys if needed
+            self.q_norm = (
+                args.norm_type(
+                    args.n_heads * self.head_dim,
+                    eps=args.norm_eps,
+                )
+                if self.apply_qk_norm
+                else nn.Identity()
+            )
+            self.k_norm = (
+                args.norm_type(
+                    args.n_heads * self.head_dim,
+                    eps=args.norm_eps,
+                )
+                if self.apply_qk_norm
+                else nn.Identity()
             )
-            if self.apply_qk_norm
-            else nn.Identity()
-        )
 
         self.layer_id = layer_id
         self.dim = args.dim
         self.reset_parameters()
 
     def reset_parameters(self):
-        # initialize weights by trunc_normal(1/sqrt(fan_in))
-        std = 1.0 / math.sqrt(self.dim)
-        torch.nn.init.trunc_normal_(self.in_proj.weight, std=std, a=-3 * std, b=3 * std)
-        # scale init by depth as in https://arxiv.org/abs/1908.11365 -- worked slightly better.
-        std = std / math.sqrt(2 * (self.layer_id + 1))
-        torch.nn.init.trunc_normal_(self.out_proj.weight, std=std, a=-3 * std, b=3 * std)
+        if using_te:
+            # initialize weights by trunc_normal(1/sqrt(fan_in))
+            std = 1.0 / math.sqrt(self.dim)
+            torch.nn.init.trunc_normal_(self.in_proj.weight_tensor.float(), std=std, a=-3 * std, b=3 * std)
+            # scale init by depth as in https://arxiv.org/abs/1908.11365 -- worked slightly better.
+            std = std / math.sqrt(2 * (self.layer_id + 1))
+            torch.nn.init.trunc_normal_(self.out_proj.weight_tensor.float(), std=std, a=-3 * std, b=3 * std)
+        else:
+            # initialize weights by trunc_normal(1/sqrt(fan_in))
+            std = 1.0 / math.sqrt(self.dim)
+            torch.nn.init.trunc_normal_(self.in_proj.weight, std=std, a=-3 * std, b=3 * std)
+            # scale init by depth as in https://arxiv.org/abs/1908.11365 -- worked slightly better.
+            std = std / math.sqrt(2 * (self.layer_id + 1))
+            torch.nn.init.trunc_normal_(self.out_proj.weight, std=std, a=-3 * std, b=3 * std)
 
     def forward(self, x: torch.Tensor, is_causal=True, past_key_value=None, use_cache=False):
         batchsize, q_len, _ = x.shape
@@ -202,9 +245,14 @@ def __init__(self, layer_id, args: Params):
         elif args.ffn_type == "gelu":
             # Follows mosaic mpt7b, but without a bias.
             self.hidden_dim = args.dim * 4
-            self._ff_w1 = nn.Linear(args.dim, self.hidden_dim, bias=False)
-            self._ff_w2 = nn.Linear(self.hidden_dim, args.dim, bias=False)
-            self.feed_forward = nn.Sequential(self._ff_w1, nn.GELU(approximate="none"), self._ff_w2)
+            if using_te:
+                self._ff_w1 = te.Linear(args.dim, self.hidden_dim, bias=False, device="cuda")
+                self._ff_w2 = te.Linear(self.hidden_dim, self.dim, bias=False, device="cuda")
+                self.feed_forward = nn.Sequential(self._ff_w1, nn.GELU(approximate="none"), self._ff_w2)
+            else:
+                self._ff_w1 = nn.Linear(args.dim, self.hidden_dim, bias=False)
+                self._ff_w2 = nn.Linear(self.hidden_dim, args.dim, bias=False)
+                self.feed_forward = nn.Sequential(self._ff_w1, nn.GELU(approximate="none"), self._ff_w2)
         elif args.ffn_type == "moe":
             moe_args = MoEArgs(
                 hidden_size=args.dim,
@@ -222,14 +270,24 @@ def __init__(self, layer_id, args: Params):
             self.feed_forward = MoE(moe_args)
 
         self.layer_id = layer_id
-        self.attention_norm = args.norm_type(
-            args.dim,
-            eps=args.norm_eps,
-        )
-        self.ffn_norm = args.norm_type(
-            args.dim,
-            eps=args.norm_eps,
-        )
+        if using_te:
+            self.attention_norm = te.LayerNorm(
+                args.dim,
+                eps=args.norm_eps,
+            )
+            self.ffn_norm = te.LayerNorm(
+                args.dim,
+                eps=args.norm_eps,
+            )
+        else:
+            self.attention_norm = args.norm_type(
+                args.dim,
+                eps=args.norm_eps,
+            )
+            self.ffn_norm = args.norm_type(
+                args.dim,
+                eps=args.norm_eps,
+            )
         self.attention.seq_len = args.seq_len
         self.reset_parameters()
 
@@ -243,12 +301,20 @@ def reset_parameters(self):
             std = std / math.sqrt(2 * (self.layer_id + 1))
             torch.nn.init.trunc_normal_(self.feed_forward.w3.weight, std=std, a=-3 * std, b=3 * std)
         elif self._ffn_type == "gelu":
-            std = 1.0 / math.sqrt(self.dim)
-            torch.nn.init.trunc_normal_(self._ff_w1.weight, std=std, a=-3 * std, b=3 * std)
+            if using_te:
+                std = 1.0 / math.sqrt(self.dim)
+                torch.nn.init.trunc_normal_(self._ff_w1.weight_tensor.float(), std=std, a=-3 * std, b=3 * std)
 
-            std = 1.0 / math.sqrt(self.hidden_dim)
-            std = std / math.sqrt(2 * (self._layer_id + 1))
-            torch.nn.init.trunc_normal_(self._ff_w2.weight, std=std, a=-3 * std, b=3 * std)
+                std = 1.0 / math.sqrt(self.hidden_dim)
+                std = std / math.sqrt(2 * (self._layer_id + 1))
+                torch.nn.init.trunc_normal_(self._ff_w2.weight_tensor.float(), std=std, a=-3 * std, b=3 * std)
+            else:
+                std = 1.0 / math.sqrt(self.dim)
+                torch.nn.init.trunc_normal_(self._ff_w1.weight, std=std, a=-3 * std, b=3 * std)
+
+                std = 1.0 / math.sqrt(self.hidden_dim)
+                std = std / math.sqrt(2 * (self._layer_id + 1))
+                torch.nn.init.trunc_normal_(self._ff_w2.weight, std=std, a=-3 * std, b=3 * std)
 
     def forward(self, x, past_key_value=None, use_cache=False):
         h, past_key_value = self.attention(

open_lm/norms.py

@@ -8,6 +8,18 @@
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
 
+# Adding flag if using TE FP8
+using_te = False
+try:
+    import transformer_engine.pytorch as te
+    from transformer_engine.common import recipe
+
+    fp8_format = recipe.Format.HYBRID
+    fp8_recipe = recipe.DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max")
+    using_te = True
+except ImportError as ie:
+    using_te = False
+
 
 class LayerNorm(nn.Module):
     # NOTE: taken from official pytorch implementation and modified
@@ -55,7 +67,16 @@ def reset_parameters(self) -> None:
                 self.bias.zero_()
 
     def forward(self, input: Tensor) -> Tensor:
-        return F.layer_norm(input, self.normalized_shape, self.weight, self.bias, self.eps)
+        if using_te:
+            layer_norm_module = te.LayerNorm(
+                self.normalized_shape, eps=self.eps, device="cuda", params_dtype=input.dtype
+            )
+            output_tensor = layer_norm_module(input)
+            if self.weight is not None and self.bias is not None:
+                output_tensor = output_tensor * self.weight + self.bias
+            return output_tensor
+        else:
+            return F.layer_norm(input, self.normalized_shape, self.weight, self.bias, self.eps)
 
     def extra_repr(self) -> str:
         return (
@@ -77,13 +98,22 @@ def forward(self, x):
         downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
         downcast_bias = _cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias
         with torch.autocast(enabled=False, device_type=module_device.type):
-            return F.layer_norm(
-                downcast_x,
-                self.normalized_shape,
-                downcast_weight,
-                downcast_bias,
-                self.eps,
-            )
+            if using_te:
+                layer_norm_module = te.LayerNorm(
+                    self.normalized_shape, eps=self.eps, device="cuda", params_dtype=downcast_x.dtype
+                )
+                output_tensor = layer_norm_module(downcast_x)
+                if downcast_weight is not None and downcast_bias is not None:
+                    output_tensor = output_tensor * downcast_weight + downcast_bias
+                return output_tensor
+            else:
+                return F.layer_norm(
+                    downcast_x,
+                    self.normalized_shape,
+                    downcast_weight,
+                    downcast_bias,
+                    self.eps,
+                )
 
 
 def _cast_if_autocast_enabled(tensor):

open_lm/params.py

@@ -742,6 +742,17 @@ def parse_args(args):
         action="store_true",
         help="If set, allow model to do multiple data passes over our dataset, in order to reach the desired number of tokens.",
     )
+    parser.add_argument(
+        "--use-smp-flash-attention",
+        type=int,
+        default=None,
+        help="Using SMP Flash Attention.",
+    )
+    parser.add_argument(
+        "--sharding-strategy",
+        default=None,
+        help="Sharding Strategy",
+    )
 
     add_model_args(parser)