Fix and test LM head (#240)

Author: jlamypoirier

fast_llm/engine/multi_stage/fsdp.py

@@ -167,7 +167,8 @@ def setup(
         grad_shard: torch.Tensor | None,
         weight_buffer: torch.Tensor | None,
         grad_buffer: torch.Tensor | None,
-        sequence_tensor_parallel: bool = False,
+        sequence_tensor_parallel: bool,
+        device: torch.device | None,
     ) -> None:
         assert not self._is_setup
         self._is_setup = True
@@ -176,11 +177,19 @@ def setup(
 
         # Validate and set the shards and buffers
         if self._mode.on_device:
-            self._weight_shard = self._weight_shard_meta.validate(weight_shard)
+            self._weight_shard = (
+                torch.empty_like(self._weight_shard_meta, device=device)
+                if weight_shard is None
+                else self._weight_shard_meta.validate(weight_shard)
+            )
         else:
             Assert.none(weight_shard)
         if self._mode.support_forward:
-            self._weight_buffer = self._weight_buffer_meta.validate(weight_buffer)
+            self._weight_buffer = (
+                torch.empty_like(self._weight_buffer_meta, device=device)
+                if weight_buffer is None
+                else self._weight_buffer_meta.validate(weight_buffer)
+            )
             # Pre-compute the local shard for restore ops.
             self._weight_buffer_local_shard = self._weight_buffer[
                 self._fsdp_dim.rank * self._shard_size : (self._fsdp_dim.rank + 1) * self._shard_size
@@ -189,8 +198,16 @@ def setup(
             Assert.none(weight_buffer)
 
         if self._mode.support_backward:
-            self._grad_shard = self._grad_shard_meta.validate(grad_shard)
-            self._grad_buffer = self._grad_buffer_meta.validate(grad_buffer)
+            self._grad_shard = (
+                torch.empty_like(self._grad_shard_meta, device=device)
+                if grad_shard is None
+                else self._grad_shard_meta.validate(grad_shard)
+            )
+            self._grad_buffer = (
+                torch.empty_like(self._grad_buffer_meta, device=device)
+                if grad_buffer is None
+                else self._grad_buffer_meta.validate(grad_buffer)
+            )
             # Pre-compute the local shard for reduce ops.
             self._grad_buffer_local_shard = self._grad_buffer[
                 self._fsdp_dim.rank * self._shard_size : (self._fsdp_dim.rank + 1) * self._shard_size

fast_llm/engine/multi_stage/stage.py

@@ -1,3 +1,4 @@
+import collections
 import logging
 import typing
 
@@ -38,13 +39,13 @@ def setup(  # noqa
         self,
         *,
         distributed: Distributed,
-        weight_shards: list[torch.Tensor | None] | None,
-        grad_shards: list[torch.Tensor | None] | None,
-        weight_buffers: list[torch.Tensor | None] | None,
-        grad_buffers: list[torch.Tensor | None] | None,
+        weight_shards: list[torch.Tensor | None] | None = None,
+        grad_shards: list[torch.Tensor | None] | None = None,
+        weight_buffers: list[torch.Tensor | None] | None = None,
+        grad_buffers: list[torch.Tensor | None] | None = None,
         mode: StageMode = StageMode.training,
         is_tied_weight_copy: bool = False,
-        weight_buffer_shared_with: list["Stage"],
+        weight_buffer_shared_with: collections.abc.Sequence["Stage"] = (),
     ) -> None:
         super().setup(
             distributed=distributed,
@@ -92,7 +93,11 @@ def forward_meta(self, input_: TensorMeta, kwargs: dict) -> TensorMeta:
         return input_
 
     def forward(
-        self, input_: torch.Tensor, kwargs: dict, losses: dict[str, list[torch.Tensor]], metrics: dict | None = None
+        self,
+        input_: torch.Tensor,
+        kwargs: dict,
+        losses: dict[str, list[torch.Tensor]] | None = None,
+        metrics: dict | None = None,
     ) -> tuple[torch.Tensor | None, tuple[torch.Tensor | None, torch.Tensor | None]]:
         assert self._is_restored
         assert self._mode.support_forward

fast_llm/engine/multi_stage/stage_base.py

@@ -6,7 +6,7 @@
 
 from fast_llm.config import Configurable
 from fast_llm.core.distributed import check_parallel_match
-from fast_llm.engine.base_model.base_model import BaseModel
+from fast_llm.engine.base_model.base_model import BaseModel, Layer
 from fast_llm.engine.config_utils.data_type import DataType
 from fast_llm.engine.distributed.config import DistributedConfig, DistributedDimNames
 from fast_llm.engine.distributed.distributed import Distributed
@@ -29,7 +29,7 @@ def __init__(
         self,
         *,
         config: StageConfig,
-        base_model: BaseModel,
+        base_model: BaseModel | list[Layer],
         distributed_config: DistributedConfig,
         begin: int,
         end: int,
@@ -153,6 +153,7 @@ def setup(
                 weight_buffer=weight_buffer,
                 grad_buffer=grad_buffer,
                 sequence_tensor_parallel=self._distributed_config.sequence_tensor_parallel,
+                device=self._distributed.device,
             )
 
         if self._mode.support_forward:

fast_llm/functional/triton/normalization.py

@@ -176,14 +176,17 @@ def triton_normalization_forward(
     training: bool,
     zero_centered: bool,
 ) -> tuple[torch.Tensor, list[typing.Any]] | None:
+    # Note: Converting input automatically to training dtype to match Apex behaviour,
+    #  needed for full precision residual.
+    # TODO: Review this?
     assert weight.shape == input_.shape[-1:]
     if bias is not None:
         assert weight.shape == bias.shape
     assert input_.is_contiguous()
     n_rows = input_.shape[:-1].numel()
     n_cols = weight.numel()
 
-    output = torch.empty_like(input_)
+    output = torch.empty_like(input_, dtype=weight.dtype)
     inv_var = torch.empty(n_rows, dtype=torch.float32, device="cuda")
 
     block_size = triton.next_power_of_2(n_cols)

fast_llm/layers/common/auxiliary_loss.py

@@ -16,7 +16,7 @@ def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor | None, ...]:
 def calculate_z_loss(logits: torch.Tensor, logits_scale_factor: float = 1.0) -> torch.Tensor:
     if logits_scale_factor != 1.0:
         logits *= logits_scale_factor
-    return torch.mean(torch.square(torch.logsumexp(logits, dim=-1)))
+    return torch.mean(torch.logsumexp(logits, dim=-1) ** 2)
 
 
 def z_loss(

fast_llm/layers/common/normalization.py

@@ -3,7 +3,7 @@
 from fast_llm.engine.config_utils.run import log_main_rank
 from fast_llm.engine.config_utils.tensor_space import TensorDim
 from fast_llm.functional.config import TritonConfig
-from fast_llm.functional.triton.normalization import rms_norm, triton_normalization_autograd
+from fast_llm.functional.triton.normalization import triton_normalization_autograd
 from fast_llm.layers.common.config import NormalizationImplementation
 from fast_llm.tensor import ParameterMeta, accumulate_gradient, init_ones_, init_zeros_
 from fast_llm.utils import Assert
@@ -141,6 +141,9 @@ def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None,
 class LayerNorm(torch.nn.Module):
     """
     A layer normalization layer, supporting multiple implementations.
+    Note: Converting input automatically to training dtype to match Apex behaviour,
+     needed for full precision residual.
+    TODO: Review this?
     """
 
     def __init__(
@@ -214,12 +217,15 @@ def _forward_fused(self, input_: torch.Tensor) -> torch.Tensor:
         return FusedLayerNorm.apply(input_, self.normalized_shape, self.weight, self.bias, self._eps)
 
     def _forward_torch(self, input_: torch.Tensor) -> torch.Tensor:
-        return torch.nn.functional.layer_norm(input_, self.normalized_shape, self.weight, self.bias, self._eps)
+        return torch.layer_norm(input_.to(self.weight.dtype), self.normalized_shape, self.weight, self.bias, self._eps)
 
 
 class RMSNorm(torch.nn.Module):
     """
     A RMS normalization layer.
+    Note: Converting input automatically to training dtype to match Apex behaviour,
+     needed for full precision residual.
+    TODO: Review this?
     """
 
     def __init__(
@@ -276,4 +282,4 @@ def _forward_fused(self, input_: torch.Tensor) -> torch.Tensor:
         return FusedRMSNorm.apply(input_, self.normalized_shape, self.weight, self._eps)
 
     def _forward_torch(self, input_: torch.Tensor) -> torch.Tensor:
-        return rms_norm(input_, self.weight, self._eps)
+        return torch.rms_norm(input_.to(self.weight.dtype), self.normalized_shape, self.weight, self._eps)

fast_llm/layers/language_model/head.py

@@ -50,7 +50,9 @@ def __init__(
         self._group_size = tensor_space.distributed_config.tensor_parallel
         self._sequence_parallel = tensor_space.distributed_config.sequence_tensor_parallel
         self._parallel_embeddings = tensor_space.distributed_config.tensor_parallel > 1 and config.parallel_embeddings
-        self._sequence_parallel_logits = self._sequence_parallel and not self._parallel_embeddings
+        self._sequence_parallel_logits = (
+            tensor_space.distributed_config.sequence_tensor_parallel and not config.parallel_embeddings
+        )
         self._cross_entropy_splits = config.cross_entropy_splits
         if self._cross_entropy_splits is not None and self._sequence_parallel:
             assert not self._parallel_embeddings
@@ -67,7 +69,7 @@ def __init__(
         # >0: multi-token prediction (MTP)
         Assert.geq(prediction_distance, 0)
         self._prediction_distance = prediction_distance
-        self.is_last_head = self._prediction_distance == config.prediction_heads - 1
+        self._is_last_head = self._prediction_distance == config.prediction_heads - 1
 
         self._init_output_weights(hidden_dim, config)
 
@@ -114,7 +116,7 @@ def forward(
                 tensor_name="Loss",
                 reductions=((DistributedDimNames.data, ReduceOp.AVG),),  # noqa
             )
-        if not self.is_last_head:
+        if not self._is_last_head:
             # MTP: split the stacked input
             shared_hidden, input_ = torch.unbind(input_, dim=0)
         # TODO: Pytorch copies the grads in backward for no reason (not sure if still the case)
@@ -123,10 +125,10 @@ def forward(
         # TODO: Drop autograd entirely.
         # TODO: Skip cross-entropy backward if not needed.
         language_model_loss = self._forward(input_, kwargs, losses)
-        if language_model_loss is not None:
+        if losses is not None and language_model_loss is not None:
             losses[self._loss_name].append(language_model_loss)
         # TODO: Return the model output when needed.
-        if self.is_last_head:
+        if self._is_last_head:
             # Last head should return the loss for backward.
             return language_model_loss
         else:
@@ -147,14 +149,13 @@ def _forward_backward(
         if target is not None:
             if self._config.distillation_model is None:
                 # MTP: Shift the labels
-                target = (
-                    target[self._prediction_distance : self._prediction_distance + input_.size(0),]
-                    if kwargs[TransformerKwargs.sequence_first]
-                    else target[
-                        :,
-                        self._prediction_distance : self._prediction_distance + input_.size(1),
-                    ]
+                target_sequence_length = (
+                    target.size(1 - kwargs[TransformerKwargs.sequence_first]) + 1 - self._config.prediction_heads
                 )
+                if TransformerKwargs.sequence_q_dim in kwargs:
+                    Assert.eq(target_sequence_length, kwargs[TransformerKwargs.sequence_q_dim].size)
+                target_slice = slice(self._prediction_distance, self._prediction_distance + target_sequence_length)
+                target = target[target_slice] if kwargs[TransformerKwargs.sequence_first] else target[:, target_slice]
                 target = target.flatten()
             else:
                 # Target is reference model logits.

fast_llm/layers/transformer/attention.py

@@ -84,7 +84,7 @@ def __init__(
         super().__init__()
         self._config = config
         self._tensor_space = tensor_space
-        Assert.in_range_incl(layer_index, 1, self._config.num_layers)
+        Assert.in_range_incl(layer_index, 1, max(self._config.num_layers, 1))
         self._layer_index = layer_index
         self._sequence_parallel = self._tensor_space.distributed_config.sequence_tensor_parallel
         self._debug_transformer = self._config.debug_transformer

fast_llm/layers/transformer/config.py

@@ -674,11 +674,11 @@ def _validate(self) -> None:
             if self.init_method_std_qkv is None:
                 self.init_method_std_qkv = self.init_method_std
             if self.init_method_std_attn_proj is None:
-                self.init_method_std_attn_proj = self.init_method_std / (2 * self.num_layers) ** 0.5
+                self.init_method_std_attn_proj = self.init_method_std / max(2 * self.num_layers, 1) ** 0.5
             if self.init_method_std_mlp_1 is None:
                 self.init_method_std_mlp_1 = self.init_method_std
             if self.init_method_std_mlp_2 is None:
-                self.init_method_std_mlp_2 = self.init_method_std / (2 * self.num_layers) ** 0.5
+                self.init_method_std_mlp_2 = self.init_method_std / max(2 * self.num_layers, 1) ** 0.5
             if self.init_method_max_qkv is None:
                 self.init_method_max_qkv = self.init_method_max
             if self.init_method_min_qkv is None:

fast_llm/models/gpt/model.py

@@ -72,34 +72,30 @@ def __init__(
             self._preprocessors.append(BackupAttentionPreprocessor(self._config.transformer, self._tensor_space))
 
     def get_output_layers(self) -> list[Layer]:
-        return [
-            layer
-            for i in range(self._config.prediction_heads)
-            for layer in [
-                TransformerLayer(
-                    self._config.transformer,
-                    self._tensor_space,
-                    # TODO MTP: which index?
-                    layer_index=self._config.transformer.num_layers,
-                    # The last layer only returns the transformer output.
-                    # The previous layers return a stack of shared_hidden and transformer_output.
-                    return_input=i < self._config.prediction_heads - 1,
-                ),
+        layers = []
+        for i in range(self._config.prediction_heads):
+            if i > 0:
+                layers.append(
+                    TransformerLayer(
+                        self._config.transformer,
+                        self._tensor_space,
+                        # TODO MTP: which index?
+                        layer_index=max(self._config.transformer.num_layers, 1),
+                        # The last layer only returns the transformer output.
+                        # The previous layers return a stack of shared_hidden and transformer_output.
+                        return_input=i < self._config.prediction_heads - 1,
+                    )
+                )
+            layers.append(
                 LanguageModelHead(
                     self._config,
                     self._tensor_space,
                     prediction_distance=i,
-                ),
-            ]
-        ]
+                )
+            )
+        return layers
 
     def get_layers(self) -> list[Layer]:
-        if self._config.transformer.num_layers == 0:
-            Assert.eq(self._config.prediction_heads, 1)
-            return [
-                LanguageModelEmbedding(self._config, self._tensor_space),
-                LanguageModelHead(self._config, self._tensor_space, 0),
-            ]
         return [
             LanguageModelEmbedding(self._config, self._tensor_space),
             *[
@@ -108,7 +104,7 @@ def get_layers(self) -> list[Layer]:
                     self._tensor_space,
                     layer_index=i + 1,
                 )
-                for i in range(self._config.transformer.num_layers - 1)
+                for i in range(self._config.transformer.num_layers)
             ],
             *self.get_output_layers(),
         ]