Mamba2 to be merged

fast_llm/engine/config_utils/tensor_space.py

@@ -66,13 +66,23 @@ def replace_parallel_dim(self, distributed_dim: DistributedDim) -> typing.Self:
         return TensorDim(self.name, self.size * distributed_dim.size, distributed_dim)
 
     def local_to_global(self, tensor: "torch.Tensor", dim: int = 0) -> "torch.Tensor":
-        if self.parallel_group is not None:
+        if self.is_parallel:
             from fast_llm.core.ops import gather_op
 
             return gather_op(tensor, self.parallel_group, dim)
         else:
             return tensor
 
+    def local_to_global_partial(
+        self, tensor: "torch.Tensor", dim: int = 0, fill_value: float | int = -1
+    ) -> "torch.Tensor":
+        if self.is_parallel:
+            output = tensor.new_full((*tensor.shape[:dim], self.parallel_dim.size, *tensor.shape[dim:]), fill_value)
+            output.narrow(dim, self.parallel_dim.rank, 1).copy_(tensor.unsqueeze(dim)).squeeze(dim)
+            return output.flatten(dim, dim + 1)
+        else:
+            return tensor
+
     def global_to_local(self, tensor: "torch.Tensor", dim: int = 0, expand: bool = False) -> "torch.Tensor":
         return (
             tensor.chunk(self.parallel_dim.size, dim)[self.parallel_dim.rank]
@@ -85,7 +95,7 @@ class CompositeTensorDim(TensorDim):
     def __init__(self, name: str, tensor_dims: tuple[TensorDim, ...]):
         parallel_dim = None
         for dim, tensor_dim in enumerate(tensor_dims):
-            if tensor_dim.is_parallel:
+            if tensor_dim.parallel_dim is not None:
                 # TODO: Allow more than one parallel subdim?
                 assert parallel_dim is None
                 parallel_dim = tensor_dim.parallel_dim
@@ -111,6 +121,15 @@ def local_to_global(self, tensor: "torch.Tensor", dim: int = 0) -> "torch.Tensor
 
         return tensor.flatten(dim, dim + len(self._tensor_dims) - 1)
 
+    def local_to_global_partial(
+        self, tensor: "torch.Tensor", dim: int = 0, fill_value: float | int = -1
+    ) -> "torch.Tensor":
+        tensor = tensor.unflatten(dim, [tensor_dim.size for tensor_dim in self._tensor_dims])
+        for i, tensor_dim in enumerate(self._tensor_dims):
+            tensor = tensor_dim.local_to_global_partial(tensor, dim + i)
+
+        return tensor.flatten(dim, dim + len(self._tensor_dims) - 1)
+
     def global_to_local(self, tensor: "torch.Tensor", dim: int = 0, expand: bool = False) -> "torch.Tensor":
         tensor = tensor.unflatten(dim, [tensor_dim.global_size for tensor_dim in self._tensor_dims])
         for i, tensor_dim in reversed(list(enumerate(self._tensor_dims))):
@@ -157,6 +176,27 @@ def local_to_global(self, tensor: "torch.Tensor", dim: int = 0) -> "torch.Tensor
             else tensor
         )
 
+    def local_to_global_partial(
+        self, tensor: "torch.Tensor", dim: int = 0, fill_value: float | int = -1
+    ) -> "torch.Tensor":
+        import torch
+
+        return (
+            torch.concatenate(
+                [
+                    tensor_dim.local_to_global_partial(tensor_, dim)
+                    for tensor_, tensor_dim in zip(
+                        tensor.split([tensor_dim.size for tensor_dim in self._tensor_dims], dim),
+                        self._tensor_dims,
+                        strict=True,
+                    )
+                ],
+                dim,
+            )
+            if self.is_parallel
+            else tensor
+        )
+
     def global_to_local(self, tensor: "torch.Tensor", dim: int = 0, expand: bool = False) -> "torch.Tensor":
         if self.is_parallel and expand:
             raise NotImplementedError()
@@ -223,8 +263,5 @@ def add_tensor_dim(self, tensor_dim: TensorDim) -> None:
                 )
             self._tensor_dims[tensor_dim.name] = tensor_dim
 
-    def get_tensor_dim(self, name: str) -> TensorDim:
+    def __getitem__(self, name: str) -> TensorDim:
         return self._tensor_dims[name]
-
-    # TODO: Replace uses
-    __getitem__ = get_tensor_dim

fast_llm/engine/multi_stage/config.py

@@ -31,6 +31,7 @@
 
 if typing.TYPE_CHECKING:
     from fast_llm.engine.inference.huggingface import HuggingfaceBaseModelForCausalLM
+    from fast_llm.engine.inference.runner import InferenceRunner
     from fast_llm.engine.multi_stage.fast_llm_model import FastLLMModel
 
 logger = logging.getLogger(__name__)
@@ -241,6 +242,10 @@ def get_checkpoint_handler_class(cls, format: type[CheckpointFormat] | str) -> t
     def get_model_class(cls) -> type["FastLLMModel"]:
         raise NotImplementedError
 
+    @classmethod
+    def get_inference_runner_class(cls) -> type["InferenceRunner"]:
+        raise NotImplementedError
+
     @classmethod
     def get_huggingface_model_for_causal_lm_class(cls) -> type["HuggingfaceBaseModelForCausalLM"]:
         raise NotImplementedError

fast_llm/engine/multi_stage/fsdp.py

@@ -441,39 +441,21 @@ def _get_parameter_shard_indices_in_full_weight(
         where it is located in the shard if it exists, or -1 if it's not in the shard.
         Used to determine the location of each entry in a different distributed configuration.
         """
-
-        # Create an empty index for the global parameter.
-        index = torch.full(
-            parameter_meta.global_shape,
-            -1,
-            dtype=torch.int64,
-            device=device,
-        )
         # Set the shard slice of the global parameter to corresponding indices of the parameter slice of the shard
         begin, end = self._get_parameter_range_in_shard(parameter_name)
 
-        buffer_index = parameter_meta.global_to_local(index, expand=True)
-        # Copying directly into `buffer_index` requires a view of the tensor, which may not be feasible.
-        # In that case, we work with a separate tensor to be copied back into `buffer_index`.
-        try:
-            buffer_index_flat = buffer_index.view(-1)
-            is_view = True
-        except RuntimeError:
-            buffer_index_flat = buffer_index.new_full((buffer_index.numel(),), -1)
-            is_view = False
-
-        # Copy the shard indices at their respective positions in the flat buffer index.
-        buffer_index_flat[
+        # Create an empty local index to hold the local shard indices.
+        buffer_index = torch.full_like(parameter_meta, -1, dtype=torch.int64, device=device)
+
+        # Copy the shard indices at their respective positions in the buffer index.
+        buffer_index.flatten()[
             self._index_buffer_to_param(
                 self._fsdp_dim.rank * self._shard_size, parameter_name
             ) : self._index_buffer_to_param((self._fsdp_dim.rank + 1) * self._shard_size, parameter_name)
         ].copy_(torch.arange(begin, end, dtype=torch.int64, device=device))
 
-        # If needed, copy the flat buffer index back into the index.
-        if not is_view:
-            buffer_index.copy_(buffer_index_flat.view_as(buffer_index))
-
-        return index
+        # Create a global index from the local one.
+        return parameter_meta.local_to_global_partial(buffer_index, -1)
 
     def copy_shard_overlaps(
         self,

fast_llm/engine/training/config.py

@@ -32,7 +32,6 @@
 from fast_llm.utils import Assert
 
 if typing.TYPE_CHECKING:
-    from fast_llm.engine.inference.runner import InferenceRunner
     from fast_llm.engine.training.trainer import Trainer, TrainingEvaluator
 
 
@@ -403,10 +402,6 @@ def _setup(self):
     def get_trainer_class(cls) -> type["Trainer"]:
         raise NotImplementedError
 
-    @classmethod
-    def get_inference_runner_class(cls) -> type["InferenceRunner"]:
-        raise NotImplementedError
-
     def _get_runnable(self) -> typing.Callable[[], None]:
         from fast_llm.engine.distributed.distributed import Distributed
 

fast_llm/engine/training/trainer.py

@@ -142,7 +142,7 @@ def __init__(self, config: TrainerConfig):
         self._reference_models = {}
         for name, reference_config in self._config.reference_models.items():
             log_main_rank(f"Creating `{name} reference model...")
-            self._reference_models[name] = self._config.get_inference_runner_class()(
+            self._reference_models[name] = reference_config.model.get_inference_runner_class()(
                 reference_config.model.get_model_class()(reference_config.model)
             )
             self._multi_stage.base_model.add_reference_model(name, self._reference_models[name])

fast_llm/functional/cross_entropy.py

@@ -49,6 +49,19 @@ def _torch_cross_entropy_forward_backward(
     return loss.detach_(), grad
 
 
+def distributed_log_softmax(logits: torch.Tensor, group: ProcessGroup, dim: int = -1):
+    logits = logits.float()
+    local_max = logits.max(dim=dim, keepdim=True)[0]
+    all_reduce(local_max, op=ReduceOp.MAX, group=group)
+
+    logits_shifted = logits - local_max
+    exp_logits = torch.exp(logits_shifted)
+    sum_exp = exp_logits.sum(dim=dim, keepdim=True)
+    all_reduce(sum_exp, op=ReduceOp.SUM, group=group)
+
+    return logits_shifted - sum_exp.log()  # log_softmax
+
+
 @torch.compile
 def _fused_softmax_base(
     logits: torch.Tensor, logits_scale_factor: float = 1.0, group: ProcessGroup | None = None, dim: int = -1
@@ -214,38 +227,88 @@ def cross_entropy_forward_backward(
         )
 
 
-def _torch_reverse_kl_forward_backward(
+def _torch_reverse_kl_forward_backward_vocab_parallel(
     logits: torch.Tensor,
     target: torch.Tensor,
     loss_mask: torch.Tensor | None,
     grad_output: float | None,
-    logits_scale_factor: float,
     target_format: TargetFormat,
     group: ProcessGroup | None = None,
-    teacher_softmax_temperature: float = 1.0,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
     """
     Reverse KL using PyTorch's native kl_div function.
+    This is used for TP version where we split accross vocab dimantion.
     This works with sequence-tensor-parallel (distributing over the sequence dimention) as well as a non-TP case.
     In sequence-tensor-parallel, where we split along sequence dim., we compute per split loss and then average the loss.
     """
+    # TODO: merge into single function _torch_reverse_kl_forward_backward
     Assert.eq(target_format, TargetFormat.logits, msg="Reverse KL only supports logits format")
     Assert.eq(target.shape, logits.shape)
     assert target.dtype.is_floating_point, target.dtype
     if loss_mask is not None:
         Assert.eq(loss_mask.shape, logits.shape[:-1])
 
     # Compute log probabilities - let _fused_softmax handle scaling internally
-    # teacher_probs = _fused_softmax(target, logits_scale_factor * (1 / teacher_softmax_temperature), group)
-    # # teacher_log_probs = torch.log(teacher_probs + 1e-8)  # log(p)
-    # teacher_probs = torch.clamp(teacher_probs, min=1e-7)  # or even 1e-6
-    # teacher_log_probs = torch.log(teacher_probs)
+    teacher_log_probs = distributed_log_softmax(target, group=group)
+    batch_size = logits.shape[0]
+    with torch.enable_grad():
+        logits_ = logits.detach().requires_grad_(grad_output is not None)
+        student_log_probs = distributed_log_softmax(logits_, group=group)
+
+        # Reverse KL: input=teacher_log_probs, target=student_probs
+        if loss_mask is None:
+            loss = torch.nn.functional.kl_div(
+                teacher_log_probs,  # input = log(p)
+                student_log_probs,  # target = log(q)
+                reduction="sum",
+                log_target=True,
+            )
+        else:
+            # Apply loss mask - this requires some reshaping
+            raise NotImplementedError("Loss mask not implemented with TP for reverse KL , it must be doublechecked")
+            loss_per_sample = torch.nn.functional.kl_div(
+                teacher_log_probs, student_log_probs, reduction="none", log_target=True
+            ).sum(dim=-1)
+            loss = (loss_per_sample * loss_mask).sum()
+
+        if group is not None and target_format != TargetFormat.labels:
+            all_reduce(loss, op=ReduceOp.SUM, group=group)
+            loss /= batch_size
+
+        if grad_output is not None:
+            loss.backward(torch.full_like(loss, grad_output))
+            grad = logits_.grad.to(logits.dtype)
+        else:
+            grad = None
 
+    return loss.detach_(), grad
+
+
+def _torch_reverse_kl_forward_backward(
+    logits: torch.Tensor,
+    target: torch.Tensor,
+    loss_mask: torch.Tensor | None,
+    grad_output: float | None,
+    logits_scale_factor: float,
+    target_format: TargetFormat,
+    group: ProcessGroup | None = None,
+    teacher_softmax_temperature: float = 1.0,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """
+    Reverse KL using PyTorch's native kl_div function.
+    This works with sequence-tensor-parallel (distributing over the sequence dimention) as well as a non-TP case.
+    In sequence-tensor-parallel, where we split along sequence dim., we compute per split loss and then average the loss.
+    """
+    Assert.eq(target_format, TargetFormat.logits, msg="Reverse KL only supports logits format")
+    Assert.eq(target.shape, logits.shape)
+    assert target.dtype.is_floating_point, target.dtype
+    if loss_mask is not None:
+        Assert.eq(loss_mask.shape, logits.shape[:-1])
     # Scale target logits more carefully
     scaled_target = target * (logits_scale_factor / teacher_softmax_temperature)
+    # Clamp to prevent extreme values that cause NaNs in log_softmax
+    scaled_target = torch.clamp(scaled_target, min=-100.0, max=100.0)
 
-    # Clamp to prevent extreme values before log_softmax
-    scaled_target = torch.clamp(scaled_target, min=-50, max=50)
     teacher_log_probs = torch.log_softmax(scaled_target, dim=-1)
 
     # For reverse KL: KL(q||p) = Σ q * log(q/p) = Σ q * (log(q) - log(p))
@@ -256,9 +319,10 @@ def _torch_reverse_kl_forward_backward(
         logits_ = logits.detach().requires_grad_(grad_output is not None)
 
         scaled_logits = logits_ * logits_scale_factor
-        scaled_logits = torch.clamp(scaled_logits, min=-50, max=50)
+        # Clamp to prevent extreme values that cause NaNs in log_softmax
+        scaled_logits = torch.clamp(scaled_logits, min=-100.0, max=100.0)
         student_log_probs = torch.log_softmax(scaled_logits, dim=-1)
-
+        
         # Reverse KL: input=teacher_log_probs, target=student_probs
         if loss_mask is None:
             loss = torch.nn.functional.kl_div(
@@ -279,6 +343,7 @@ def _torch_reverse_kl_forward_backward(
             loss /= group.size()
 
         if grad_output is not None:
+            # note, we never get here in TP over seq. dim.
             loss.backward(torch.full_like(loss, grad_output))
             grad = logits_.grad.to(logits.dtype)
         else:
@@ -344,6 +409,14 @@ def reverse_kl_forward_backward(
         Assert.eq(teacher_softmax_temperature, 1)
         Assert.eq(logits_scale_factor, 1)
         raise NotImplementedError("Vocab parallel reverse KL is not implemented yet.")
+        return _torch_reverse_kl_forward_backward_vocab_parallel(
+            logits,
+            target,
+            loss_mask,
+            grad_output,
+            target_format,
+            group,
+        )
     else:
         return _torch_reverse_kl_forward_backward(
             logits,

fast_llm/layers/language_model/embedding.py

@@ -46,10 +46,10 @@ def __init__(
         self._dropout_p = config.transformer.hidden_dropout
         self._use_absolute_position_embeddings = config.use_absolute_position_embeddings
 
-        hidden_dim = tensor_space.get_tensor_dim(TransformerDimNames.hidden)
-        vocab_dim = tensor_space.get_tensor_dim(
+        hidden_dim = tensor_space[TransformerDimNames.hidden]
+        vocab_dim = tensor_space[
             LanguageModelDimNames.vocab_tp if self._parallel_embeddings else LanguageModelDimNames.vocab
-        )
+        ]
 
         if self._parallel_embeddings:
             self._vocab_start_index = self._distributed_config.tensor_rank * vocab_dim.size
@@ -66,7 +66,7 @@ def __init__(
         )
         if self._use_absolute_position_embeddings:
             self.position_embeddings_weight = ParameterMeta.from_dims(
-                (tensor_space.get_tensor_dim(LanguageModelDimNames.position_embed), hidden_dim),
+                (tensor_space[LanguageModelDimNames.position_embed], hidden_dim),
                 init_method=init_normal_(
                     std=config.init_method_std_embed,
                     min_val=config.init_method_min_embed,

fast_llm/layers/language_model/head.py

@@ -61,7 +61,7 @@ def __init__(
         if self._cross_entropy_splits is not None and self._sequence_parallel:
             assert not self._parallel_embeddings
 
-        hidden_dim = self._tensor_space.get_tensor_dim(TransformerDimNames.hidden)
+        hidden_dim = self._tensor_space[TransformerDimNames.hidden]
 
         self._loss_coefficient = (
             config.prediction_loss_coefficient[prediction_distance] if config.prediction_loss_coefficient else 1.0
@@ -108,9 +108,9 @@ def _init_output_weights(self, hidden_dim: TensorDim, config) -> None:
         if self._tie_word_embeddings or self._prediction_distance > 0:
             return
         # untie embedding weights
-        vocab_dim = self._tensor_space.get_tensor_dim(
+        vocab_dim = self._tensor_space[
             LanguageModelDimNames.vocab_tp if self._parallel_embeddings else LanguageModelDimNames.vocab
-        )
+        ]
         self.output_weights = ParameterMeta.from_dims(
             (vocab_dim, hidden_dim),
             init_method=init_normal_(
@@ -237,7 +237,6 @@ def _get_targets(
                     ).flatten()
             else:
                 lm_target = None
-
         targets = (dpo_target, lm_target, distillation_target)
         # If we do distillation, no need to split it here as it has already been split in the embedding layer!
         # if we do CPT/language modeling, we need to split the targets here!
@@ -350,9 +349,9 @@ def _logits_cross_entropy_forward_backward(
                 logits_scale_factor=self._logits_scale_factor,
             )
         if self._debug_transformer and self._cross_entropy_splits is None:
-            vocab_dim = self._tensor_space.get_tensor_dim(
+            vocab_dim = self._tensor_space[
                 LanguageModelDimNames.vocab if self._sequence_parallel_logits else LanguageModelDimNames.vocab_tp
-            )
+            ]
             dims = [*kwargs[TransformerKwargs.hidden_dims][:-1], vocab_dim]
             sequence_index = 1 - int(kwargs[TransformerKwargs.sequence_first])
             dims[sequence_index] = (