fused linear cross entropy

configs/research/ladder-desync-residual/ladder-overlapped-training.yml

@@ -15,11 +15,6 @@ datasets:
 tokenizer_args:
   tokenizer_name: bigcode/starcoder
 
-# kernel_args:
-#   kernels:
-#     - cute_swiglu
-#     - cute_rmsnorm
-
 model_args:
   model_class: AutoModelForCausalLM
   pretrained_config:

dolomite_engine/enums.py

@@ -74,7 +74,9 @@ class ExperimentsTrackerName(Enum):
 
 
 class Kernel(Enum):
-    cute_rmsnorm = "cute_rmsnorm"
-    cute_swiglu_unchunked = "cute_swiglu_unchunked"
+    rmsnorm_cute = "rmsnorm_cute"
+    swiglu_unchunked_cute = "swiglu_unchunked_cute"
     mamba2_ssm = "mamba2_ssm"
     scattermoe = "scattermoe"
+    cross_entropy_cute = "cross_entropy_cute"
+    fused_linear_cross_entropy_cute = "fused_linear_cross_entropy_cute"

dolomite_engine/finetune.py

@@ -13,6 +13,7 @@
 from .distributed import dtensor_to_tensor, wrap_model_container_for_distributed_training
 from .enums import DatasetSplit, Mode, TuningMethod
 from .hf_models import disable_generation_cache
+from .kernels import enable_kernels
 from .model_wrapper import ModelWrapper, get_model_container
 from .optimization import get_optimizer_container, get_scheduler_container
 from .train_utils import all_reduce_metrics_tracker, get_torch_profiler, track_metrics
@@ -395,7 +396,7 @@ def main() -> None:
     experiments_tracker.log_args(args)
 
     # main training loop
-    with disable_generation_cache():
+    with disable_generation_cache(), enable_kernels(args.kernel_args.kernels):
         train(
             args,
             model_container=model_container,

dolomite_engine/hf_models/loss.py

@@ -7,47 +7,80 @@
 from torch.distributed.tensor.parallel import loss_parallel
 
 from ..distributed import tensor_to_dtensor
-from ..utils import ProcessGroupManager
+from ..enums import Kernel
+from ..kernels import is_kernel_allowed
+from ..utils import ProcessGroupManager, is_cute_kernels_available
+
+
+if is_cute_kernels_available():
+    from cute_kernels import cross_entropy_cute, fused_linear_cross_entropy_cute
 
 
 def get_autoregressive_language_modeling_loss(
     lm_logits: torch.Tensor,
     labels: torch.Tensor,
+    hidden_states: torch.Tensor | None = None,
+    vocab_weight: torch.Tensor | None = None,
+    logits_multiplier: float = 1,
     cu_seqlens: torch.Tensor | None = None,
     use_padding_free_transformer: bool = False,
     reduction: str = "mean",
+    shift_logits_and_labels: bool = True,
+    tensor_parallel_enabled: bool = False,
 ) -> torch.Tensor | DTensor:
-    if use_padding_free_transformer:
-        assert cu_seqlens is not None
+    if shift_logits_and_labels:
+        lm_logits = lm_logits[..., :-1, :]
+        labels = labels[..., 1:]
 
-        shift_logits = lm_logits[:-1, :]
-        shift_labels = labels[1:].to(shift_logits.device)
+    if use_padding_free_transformer:
+        if shift_logits_and_labels:
+            assert cu_seqlens is not None
 
-        # this is needed so that the last token of current example doesn't predict first token of next example
-        drop_loss_positions = cu_seqlens[1:-1] - 1
-        shift_labels[drop_loss_positions] = -100
+            # this is needed so that the last token of current example doesn't predict first token of next example
+            drop_loss_positions = cu_seqlens[1:-1] - 1
+            labels[drop_loss_positions] = -100
     else:
         assert cu_seqlens is None
 
-        # Shift so that tokens < n predict n
-        shift_logits = lm_logits[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous().to(shift_logits.device)
+    if is_kernel_allowed(Kernel.fused_linear_cross_entropy_cute):
+        assert lm_logits is None
+        assert not tensor_parallel_enabled
 
-    loss_context = nullcontext
+        loss = fused_linear_cross_entropy_cute(
+            x=hidden_states.reshape(-1, hidden_states.size(-1)),
+            weight=vocab_weight,
+            labels=labels.reshape(-1),
+            reduction=reduction,
+            logits_multiplier=logits_multiplier,
+        )
+    elif is_kernel_allowed(Kernel.cross_entropy_cute):
+        assert hidden_states is None
+        assert vocab_weight is None
+        assert not tensor_parallel_enabled
+
+        loss = cross_entropy_cute(
+            x=lm_logits.reshape(-1, lm_logits.size(-1)),
+            labels=labels.reshape(-1),
+            reduction=reduction,
+            logits_multiplier=logits_multiplier,
+        )
+    else:
+        assert logits_multiplier == 1
+        loss_context = nullcontext
 
-    if ProcessGroupManager.is_initialized() and ProcessGroupManager.is_tensor_parallel_enabled():
-        loss_context = loss_parallel
-        tp_mesh = ProcessGroupManager.get_tensor_parallel_mesh()
+        if tensor_parallel_enabled:
+            loss_context = loss_parallel
+            tp_mesh = ProcessGroupManager.get_tensor_parallel_mesh()
 
-        shift_logits = tensor_to_dtensor(shift_logits, device_mesh=tp_mesh, current_placement=Shard(-1))
-        shift_labels = tensor_to_dtensor(shift_labels, device_mesh=tp_mesh, current_placement=Replicate())
+            lm_logits = tensor_to_dtensor(lm_logits, device_mesh=tp_mesh, current_placement=Shard(-1))
+            labels = tensor_to_dtensor(labels, device_mesh=tp_mesh, current_placement=Replicate())
 
-    shift_logits = shift_logits.float()
+        lm_logits = lm_logits.float()
 
-    with loss_context():
-        loss = F.cross_entropy(
-            shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), reduction=reduction
-        )
+        with loss_context():
+            loss = F.cross_entropy(
+                input=lm_logits.reshape(-1, lm_logits.size(-1)), target=labels.reshape(-1), reduction=reduction
+            )
 
     return loss
 

dolomite_engine/hf_models/mixins/__init__.py

@@ -1,3 +1,4 @@
 from .dense import BaseModelMixin, CausalLMModelMixin, PreTrainedModelMixin
 from .dense_TP import BaseModelMixin_TP, CausalLMModelMixin_TP, PreTrainedModelMixin_TP
+from .modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from .moe_TP import CausalLMMoEModelMixin_TP

dolomite_engine/hf_models/mixins/dense/base.py

@@ -1,15 +1,14 @@
 import torch
 import torch.nn as nn
 from transformers import DynamicCache, PreTrainedModel
-from transformers.modeling_outputs import BaseModelOutputWithPast
 
 from ....utils import divide_if_divisible
 from ...cache import HybridMambaAttentionDynamicCache
 from ...config import CommonConfig
 from ...enums import PositionEmbeddingType
-from ...loss import clear_aux_loss
 from ...modeling_utils import ParameterizedEmbedding, RoPE, YaRNScaledRoPE, get_normalization_function
 from ...utils import convert_padding_free_lists_to_tensors, is_generation_cache_enabled
+from ..modeling_outputs import BaseModelOutputWithPast
 
 
 class PreTrainedModelMixin(PreTrainedModel):
@@ -196,7 +195,6 @@ def forward(
                 self._get_empty_cache(input_ids) if use_cache and past_key_values is None else past_key_values
             )
 
-        clear_aux_loss()
         mamba_mask = None
         mamba_mask_computed = False
 

dolomite_engine/hf_models/mixins/dense/main.py

@@ -1,11 +1,13 @@
 import torch
 import torch.nn.functional as F
 from transformers import DynamicCache, GenerationMixin
-from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 
+from ....enums import Kernel
+from ....kernels import is_kernel_allowed
 from ...config import CommonConfig
-from ...loss import get_autoregressive_language_modeling_loss, get_aux_loss
+from ...loss import clear_aux_loss, get_autoregressive_language_modeling_loss, get_aux_loss
 from ...modeling_utils import ParameterizedEmbedding, ParameterizedLinear
+from ..modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from .base import PreTrainedModelMixin
 
 
@@ -39,8 +41,7 @@ def set_input_embeddings(self, value: ParameterizedEmbedding) -> None:
         self.transformer.wte = value
 
     def get_output_embeddings(self) -> ParameterizedLinear:
-        if not self._tied_word_embeddings:
-            return self.lm_head
+        return self.transformer.wte if self._tied_word_embeddings else self.lm_head
 
     def set_output_embeddings(self, new_embeddings: ParameterizedLinear) -> None:
         if not self._tied_word_embeddings:
@@ -60,6 +61,8 @@ def forward(
         cu_seqlens: torch.Tensor | None = None,
         max_seqlen: torch.Tensor | None = None,
         reduction: str = "mean",
+        apply_output_projection: bool = True,
+        apply_logits_multiplier: bool = True,
     ) -> CausalLMOutputWithPast:
         assert return_dict
 
@@ -87,6 +90,8 @@ def forward(
         #     position_ids -> None or (batch_size, key_length)
         # ==========================================================================================
 
+        clear_aux_loss()
+
         transformer_outputs: BaseModelOutputWithPast = self.transformer(
             input_ids,
             past_key_values=past_key_values,
@@ -99,19 +104,38 @@ def forward(
             max_seqlen=max_seqlen,
         )
 
-        lm_logits = self.get_lm_logits(transformer_outputs.last_hidden_state)
+        lm_logits = None
+        loss = None
 
-        if self.m_width is not None:
-            lm_logits = lm_logits / self.m_width
+        if labels is None:
+            if apply_output_projection:
+                lm_logits = self.get_lm_logits(transformer_outputs.last_hidden_state)
+
+            if apply_logits_multiplier and self.m_width is not None:
+                assert apply_output_projection
+                lm_logits = lm_logits / self.m_width
+        else:
+            assert apply_output_projection
+            assert apply_logits_multiplier
+            assert not is_kernel_allowed(Kernel.fused_linear_cross_entropy_cute)
+            assert not is_kernel_allowed(Kernel.cross_entropy_cute)
+
+            lm_logits = self.get_lm_logits(transformer_outputs.last_hidden_state)
+
+            if self.m_width is not None:
+                lm_logits = lm_logits / self.m_width
 
-        loss = None
-        if labels is not None:
             loss = get_autoregressive_language_modeling_loss(
                 lm_logits=lm_logits,
                 labels=labels,
+                hidden_states=None,
+                vocab_weight=None,
+                logits_multiplier=1,
                 cu_seqlens=cu_seqlens,
                 use_padding_free_transformer=self._use_padding_free_transformer,
                 reduction=reduction,
+                shift_logits_and_labels=True,
+                tensor_parallel_enabled=False,
             )
 
         aux_loss = get_aux_loss()
@@ -125,8 +149,7 @@ def forward(
             loss=loss,
             logits=lm_logits,
             past_key_values=transformer_outputs.past_key_values,
-            hidden_states=transformer_outputs.hidden_states,
-            attentions=transformer_outputs.attentions,
+            last_hidden_state=transformer_outputs.last_hidden_state,
         )
 
     def get_lm_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:

dolomite_engine/hf_models/mixins/dense_TP/base.py

@@ -1,7 +1,6 @@
 import torch
 import torch.nn as nn
 from transformers import DynamicCache
-from transformers.modeling_outputs import BaseModelOutputWithPast
 
 from ....utils import ProcessGroupManager, divide_if_divisible
 from ...config import CommonConfig
@@ -11,6 +10,7 @@
 from ...modeling_utils_TP import Dropout_TP, Embedding_TP, get_normalization_function_TP
 from ...utils import is_generation_cache_enabled
 from ..dense import BaseModelMixin, PreTrainedModelMixin
+from ..modeling_outputs import BaseModelOutputWithPast
 
 
 class PreTrainedModelMixin_TP(PreTrainedModelMixin):

dolomite_engine/hf_models/mixins/dense_TP/main.py

@@ -3,7 +3,6 @@
 import torch
 from torch.distributed._tensor.placement_types import Replicate, Shard
 from transformers import DynamicCache
-from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 
 from ....distributed import dtensor_to_tensor, tensor_to_dtensor
 from ....utils import ProcessGroupManager, SafeTensorsWeightsManager, divide_if_divisible
@@ -12,6 +11,7 @@
 from ...loss import get_autoregressive_language_modeling_loss
 from ...modeling_utils_TP import LMHead_TP
 from ..dense import CausalLMModelMixin
+from ..modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from .base import PreTrainedModelMixin_TP
 
 
@@ -54,6 +54,8 @@ def forward(
         cu_seqlens: torch.Tensor | None = None,
         max_seqlen: torch.Tensor | None = None,
         reduction: str = "mean",
+        apply_output_projection: bool = True,
+        apply_logits_multiplier: bool = True,
     ) -> CausalLMOutputWithPast | torch.Tensor:
         assert return_dict
 
@@ -84,9 +86,12 @@ def forward(
         )
 
         if not self.is_pipeline_parallel_enabled or self.is_last_stage:
-            lm_logits = self.get_lm_logits(transformer_outputs.last_hidden_state)
+            lm_logits = None
 
-            if self.m_width is not None:
+            if apply_output_projection:
+                lm_logits = self.get_lm_logits(transformer_outputs.last_hidden_state)
+
+            if apply_logits_multiplier and self.m_width is not None:
                 lm_logits = lm_logits / self.m_width
 
         if not self.is_pipeline_parallel_enabled:
@@ -98,6 +103,8 @@ def forward(
                     cu_seqlens=cu_seqlens,
                     use_padding_free_transformer=self._use_padding_free_transformer,
                     reduction=reduction,
+                    shift_logits_and_labels=True,
+                    tensor_parallel_enabled=True,
                 )
 
         if (not self.is_pipeline_parallel_enabled or self.is_last_stage) and not output_parallel_lm_logits:
@@ -110,8 +117,7 @@ def forward(
                 loss=loss,
                 logits=lm_logits,
                 past_key_values=transformer_outputs.past_key_values,
-                hidden_states=transformer_outputs.hidden_states,
-                attentions=transformer_outputs.attentions,
+                last_hidden_state=transformer_outputs.last_hidden_state,
             )
         elif self.is_last_stage:
             output = lm_logits

dolomite_engine/hf_models/mixins/modeling_outputs.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+
+import torch
+from transformers.modeling_outputs import ModelOutput
+
+
+@dataclass
+class BaseModelOutputWithPast(ModelOutput):
+    last_hidden_state: torch.Tensor | None = None
+    past_key_values: tuple[tuple[torch.Tensor]] | None = None
+
+
+@dataclass
+class CausalLMOutputWithPast(ModelOutput):
+    loss: torch.Tensor | None = None
+    logits: torch.Tensor | None = None
+    past_key_values: tuple[tuple[torch.Tensor]] | None = None
+    last_hidden_state: torch.Tensor | None = None