Custom modeling for training (#801)

# What does this PR do?

### Custom modeling code for training

#### Features

This PR adds support for custom modeling code for training. 
Each custom modeling code can be added under
`optimum/neuron/models/training`.

Having a custom modeling allows us to implement Neuron specificities in
a cleaner way than using dynamic patching.
It becomes easy to:

- Fuse linear layers together for efficiency 
- Use custom linear layers such as `GQAQKVColumnParallelLinear`, useful
with high TP sizes.
- Use custom kernels, such as the flash attention kernel

In this PR we provide a first full custom implementation with Llama.

#### Model weight transformations

Because having a custom modeling code enables to change the vanilla
Transformers implementation, we need a way to make sure that we can load
checkpoints from Transformers, and that we can save checkpoints in the
original format as well.

To do that we provide an API with the `ModelWeightTransformationSpec`
classes.
These classes represent the transformation compared to the vanilla
Transformers implementation and are directly added in the modules
containing these transformations.

For now two exist:

- `FusedLinearsSpec`: represents a transformation when multiple linear
layers are fused into a single linear layer (possibly a parallel linear)
- `GQAQKVColumnParallelLinearSpec`: represents the transformation of
separate query, key, and value projections into a single
GQAQKVColumnParalleLinear projection.

Then during loading, saving and consolidation, we use these specs to
make sure every weight matches with Transformers weights.

#### Known issues

- There seems to be an issue when saving a checkpoint for DP > 1 during
training. After initial investigation, it seems to be a compiler bug,
but it will require more work. I suggest to work on it on a another PR.

### Training example

#### Specs

- Model: `meta-llama/Llama-3.2-3B-Instruct`
- Dataset: `databricks/databricks-dolly-15k`
- Trainer: `NeuronSFTTrainer`
- DP=4, TP=8
- Gradient accumulation steps = 16 => Effective batch size = 4 x 16 = 64
- Sequence length = 2048 with packing = True
- 3 epochs
- Learning rate = 5e-4, warmup ration = 0.3, lr scheduler type =
"cosine"

#### Loss curve

![W B Chart 24_04_2025
16_12_30](https://github.com/user-attachments/assets/cb4467cd-3e18-4560-86e0-d71d55bb08b7)


### To be done in later PRs:

- Support for PP
- Support for LoRA
- Refactor `save_pretrained` as it was done for `from_pretrained` in
this PR.
- Add test that tests overfitting

Author: michaelbenayoun

docs/source/guides/distributed_training.mdx

@@ -183,11 +183,11 @@ Just as for ZeRO-1, it is possible to wrap the optimizer class to make it lazy.
 ```python
 from torch.optim import AdamW
 from optimum.neuron import NeuronAccelerator
-from optimum.neuron.accelerate.utils import ModelParallelismPlugin
+from optimum.neuron.accelerate.utils import ModelParallelismConfig
 from optimum.neuron.distributed import lazy_load_for_parallelism
 
 tensor_parallel_size = 8
-mp_plugin = ModelParallelismPlugin(
+mp_config = ModelParallelismConfig(
     tensor_parallel_size,
     parallelize_embeddings=True,
     sequence_parallel_enabled=True,
@@ -196,7 +196,7 @@ mp_plugin = ModelParallelismPlugin(
 
 accelerator = NeuronAccelerator(
     ...
-    mp_plugin=mp_plugin,
+    mp_config=mp_config,
 )
 
 with lazy_load_for_parallelism(tensor_parallel_size=tensor_parallel_size):

docs/source/training_tutorials/sft_lora_finetune_llm.py

@@ -37,7 +37,7 @@ def training_function(script_args, training_args):
         r=16,
         lora_alpha=16,
         lora_dropout=0.05,
-        target_modules=["q_proj", "gate_proj", "v_proj", "o_proj", "k_proj", "up_proj", "down_proj"],
+        target_modules=["q_proj", "v_proj", "o_proj", "k_proj", "up_proj", "down_proj"],
         bias="none",
         task_type="CAUSAL_LM",
     )

optimum/neuron/__init__.py

@@ -70,7 +70,7 @@
         "NeuronAccelerator",
         "NeuronAcceleratorState",
         "NeuronPartialState",
-        "ModelParallelismPlugin",
+        "ModelParallelismConfig",
     ],
     "pipelines": ["pipeline"],
     "utils": ["NeuronSFTConfig", "NeuronORPOConfig", "get_peft_model"],
@@ -90,7 +90,7 @@
 _import_structure["models.yolos"] = ["NeuronYolosForObjectDetection"]
 
 if TYPE_CHECKING:
-    from .accelerate import ModelParallelismPlugin, NeuronAccelerator, NeuronAcceleratorState, NeuronPartialState
+    from .accelerate import ModelParallelismConfig, NeuronAccelerator, NeuronAcceleratorState, NeuronPartialState
     from .hf_argparser import NeuronHfArgumentParser
     from .modeling import (
         NeuronModelForAudioClassification,

optimum/neuron/accelerate/__init__.py

@@ -15,4 +15,4 @@
 
 from .accelerator import NeuronAccelerator
 from .state import NeuronAcceleratorState, NeuronPartialState
-from .utils.dataclasses import ModelParallelismPlugin, NeuronDistributedType
+from .utils.dataclasses import ModelParallelismConfig, NeuronDistributedType

optimum/neuron/accelerate/accelerator.py

@@ -16,6 +16,7 @@
 
 import collections
 import contextlib
+import inspect
 import os
 import re
 import shutil
@@ -56,7 +57,7 @@
 from .state import NeuronAcceleratorState
 from .utils import (
     AutocastBackend,
-    ModelParallelismPlugin,
+    ModelParallelismConfig,
     NeuronDistributedType,
     patch_accelerate_is_torch_xla_available,
 )
@@ -99,7 +100,7 @@ class NeuronAccelerator(Accelerator):
     def __init__(
         self,
         *args,
-        mp_plugin: Optional[ModelParallelismPlugin] = None,
+        mp_config: Optional[ModelParallelismConfig] = None,
         zero_1: bool = False,
         autocast_backend: Union[str, AutocastBackend] = "xla",
         **kwargs,
@@ -147,7 +148,7 @@ def patched_is_torch_xla_available(check_is_tpu: bool = False, check_is_gpu: boo
         accelerate.state.is_torch_xla_available = patched_is_torch_xla_available
 
         patched_accelerator_state = partial(
-            NeuronAcceleratorState, mp_plugin=mp_plugin, autocast_backend=autocast_backend
+            NeuronAcceleratorState, mp_config=mp_config, autocast_backend=autocast_backend
         )
         with Patcher([("accelerate.accelerator.AcceleratorState", patched_accelerator_state)]):
             super().__init__(**full_kwargs)
@@ -226,7 +227,7 @@ def prepare_data_loader(
                 data_loader, num_replicas=num_replicas, rank=rank, force_drop_last=force_drop_last
             )
             # No need to wrap the dataloader if we are using pipeline parallelism.
-            if use_mp_device_loader and self.state.mp_plugin.pipeline_parallel_size == 1:
+            if use_mp_device_loader and self.state.mp_config.pipeline_parallel_size == 1:
                 data_loader = MpDeviceLoader(data_loader, self.device)
         return data_loader
 
@@ -302,6 +303,14 @@ def _prepare_optimizer_for_zero_1(self, optimizer: torch.optim.Optimizer, device
 
     @patch_within_function(("accelerate.accelerator.AcceleratedOptimizer", NeuronAcceleratedOptimizer))
     def prepare_optimizer(self, optimizer: torch.optim.Optimizer, device_placement: Optional[bool] = None):
+        # If we use custom modeling, we do not have to do anything for now.
+        # We will have to do some work when supporting ZeRO-1.
+        model = self._models[0] if len(self._models) == 1 else None
+        if model is not None and inspect.getmodule(model.__class__).__name__.startswith(
+            "optimum.neuron.models.training"
+        ):
+            return super().prepare_optimizer(optimizer, device_placement=device_placement)
+
         if self.distributed_type is NeuronDistributedType.MODEL_PARALLELISM:
             optimizer = self._prepare_optimizer_for_mp(optimizer, device_placement=device_placement)
         if self.zero_1:
@@ -385,7 +394,7 @@ def _prepare_model_for_mp(
 
         tied_parameters_dict = get_tied_parameters_dict(model)
         model_main_input_name = getattr(model, "main_input_name", None)
-        model = self.state.mp_plugin.parallelize_model(model, device=self.device)
+        model = self.state.mp_config.parallelize_model(model, device=self.device)
 
         if model_main_input_name is not None:
             setattr(model, "main_input_name", model_main_input_name)
@@ -444,6 +453,16 @@ def prepare_model(
         # we get access to the model, we simply check if the flags are the best and notify the user otherwise.
         check_neuron_cc_flags_for_model(model)
 
+        if inspect.getmodule(model.__class__).__name__.startswith("optimum.neuron.models.training"):
+            # We do not want to use the cache, or output unused tensors as it would imply more communication that we do not
+            # need.
+            model.config.use_cache = False
+            model.config.output_attentions = False
+            model.config.output_hidden_states = False
+            move_model_to_device(model, self.device)
+            model = super().prepare_model(model, device_placement=False, evaluation_mode=evaluation_mode)
+            return model
+
         model = self.patch_model_for_neuron(model)
 
         if self.state.mixed_precision == "bf16":
@@ -629,9 +648,9 @@ def save_optimizer_func(accelerator, optimizer, model, output_dir, i):
                 model,
                 output_dir,
                 optimizer=optimizer,
-                use_xser=self.state.mp_plugin.use_xser,
-                async_save=self.state.mp_plugin.async_save,
-                num_local_ranks_per_step=self.state.mp_plugin.num_local_ranks_per_step,
+                use_xser=self.state.mp_config.use_xser,
+                async_save=self.state.mp_config.async_save,
+                num_local_ranks_per_step=self.state.mp_config.num_local_ranks_per_step,
             )
             logger.info(f"Parallel model and optimizer saved to the directory {output_dir}")
 

optimum/neuron/accelerate/state.py

@@ -38,7 +38,7 @@
     set_neuron_cc_flags_for_torch_amp,
 )
 from .utils import NeuronDistributedType
-from .utils.dataclasses import AutocastBackend, ModelParallelismPlugin
+from .utils.dataclasses import AutocastBackend, ModelParallelismConfig
 
 
 if is_torch_xla_available():
@@ -128,7 +128,7 @@ def __init__(
         deepspeed_plugin=None,
         fsdp_plugin=None,
         megatron_lm_plugin=None,
-        mp_plugin: Optional[ModelParallelismPlugin] = None,
+        mp_config: Optional[ModelParallelismConfig] = None,
         autocast_backend: Optional[Union[str, AutocastBackend]] = None,
         _from_accelerator: bool = False,
         **kwargs,
@@ -184,18 +184,18 @@ def __init__(
                 if mixed_precision == "bf16":
                     os.environ["NEURON_RT_STOCHASTIC_ROUNDING_EN"] = "1"
 
-                if mp_plugin is None:
-                    mp_plugin = ModelParallelismPlugin()
+                if mp_config is None:
+                    mp_config = ModelParallelismConfig()
 
-                if mp_plugin.should_parallelize:
+                if mp_config.should_parallelize:
                     self.distributed_type = NeuronDistributedType.MODEL_PARALLELISM
 
-                self.mp_plugin = mp_plugin
+                self.mp_config = mp_config
 
                 if torch.distributed.is_initialized() and not parallel_state.model_parallel_is_initialized():
                     parallel_state.initialize_model_parallel(
-                        tensor_model_parallel_size=self.mp_plugin.tensor_parallel_size,
-                        pipeline_model_parallel_size=self.mp_plugin.pipeline_parallel_size,
+                        tensor_model_parallel_size=self.mp_config.tensor_parallel_size,
+                        pipeline_model_parallel_size=self.mp_config.pipeline_parallel_size,
                     )
 
             if self.distributed_type is DistributedType.NO:

optimum/neuron/accelerate/utils/__init__.py

@@ -15,7 +15,7 @@
 
 from .dataclasses import (
     AutocastBackend,
-    ModelParallelismPlugin,
+    ModelParallelismConfig,
     NeuronDistributedType,
 )
 from .misc import patch_accelerate_is_torch_xla_available

optimum/neuron/accelerate/utils/dataclasses.py

@@ -22,6 +22,12 @@
 import torch
 
 from ...distributed import ParallelizersManager
+from ...utils import is_neuronx_distributed_available
+from ...utils.torch_xla_and_neuronx_initialization import init_process_group
+
+
+if is_neuronx_distributed_available():
+    from neuronx_distributed.parallel_layers import parallel_state
 
 
 if TYPE_CHECKING:
@@ -49,7 +55,7 @@ class AutocastBackend(str, enum.Enum):
 
 
 @dataclass
-class ModelParallelismPlugin:
+class ModelParallelismConfig:
     tensor_parallel_size: int = 1
     parallelize_embeddings: bool = True
     sequence_parallel_enabled: bool = False
@@ -62,6 +68,9 @@ class ModelParallelismPlugin:
     num_local_ranks_per_step: int = 8
     use_xser: bool = True
     async_save: bool = False
+    fuse_qkv: bool = False
+    use_flash_attention: bool = True
+    recompute_causal_mask: bool = True
 
     def __post_init__(self):
         if self.tensor_parallel_size < 1:
@@ -73,6 +82,24 @@ def __post_init__(self):
         if isinstance(self.checkpoint_dir, str):
             self.checkpoint_dir = Path(self.checkpoint_dir)
 
+        if not torch.distributed.is_initialized():
+            init_process_group()
+
+        if not parallel_state.model_parallel_is_initialized():
+            parallel_state.initialize_model_parallel(
+                tensor_model_parallel_size=self.tensor_parallel_size,
+                pipeline_model_parallel_size=self.pipeline_parallel_size,
+            )
+
+    def auto_kv_size_multiplier(self, num_key_value_heads: int) -> int:
+        kv_size_multiplier = max(1, self.tensor_parallel_size // num_key_value_heads)
+        if self.kv_size_multiplier is not None and self.kv_size_multiplier != kv_size_multiplier:
+            raise ValueError(
+                "A kv size multiplier was already specified and is different from the inferred one: "
+                f"{self.kv_size_multiplier}"
+            )
+        return kv_size_multiplier
+
     @property
     def should_parallelize(self):
         return self.tensor_parallel_size > 1 or self.pipeline_parallel_size > 1

optimum/neuron/distributed/base.py

@@ -16,6 +16,7 @@
 
 import contextlib
 import gc
+import inspect
 import math
 from abc import ABC, abstractclassmethod
 from collections import defaultdict
@@ -559,6 +560,10 @@ def parallelize(
         orig_model, peft_prefix = get_base_model_and_peft_prefix(model)
         model_class = orig_model.__class__
 
+        # We skip parallelization if the model is coming from a custom modeling since it is already parallelized.
+        if inspect.getmodule(orig_model.__class__).__name__.startswith("optimum.neuron.models.training"):
+            return orig_model
+
         if peft_prefix:
             # We update the weight_map to contain both the original parameter names, and the ones in the PeftModel.
             # The reason we keep both is because depending on the context during parallelization one or the other name

optimum/neuron/distributed/checkpointing.py

@@ -16,6 +16,7 @@
 
 import json
 import os
+from functools import partial
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Literal, Union
 
@@ -103,7 +104,7 @@ def create_gqa_query_or_output_projection_weight_from_full_weight(
     return full_weight
 
 
-def consolidate_tensor_parallel_checkpoints(
+def old_consolidate_tensor_parallel_checkpoints(
     sharded_checkpoints: List[Path],
     load_function: Callable[[Union[str, Path]], Dict[str, Any]],
     metadata: Dict[str, Any],
@@ -205,6 +206,43 @@ def consolidate_tensor_parallel_checkpoints(
     return consolidated_state_dict
 
 
+def consolidate_tensor_parallel_checkpoints(
+    sharded_checkpoints: List[Path],
+    load_function: Callable[[Union[str, Path]], Dict[str, Any]],
+    metadata: Dict[str, Any],
+) -> Dict[str, "torch.Tensor"]:
+    from ..models.training import ModelWeightTransformationSpecs, to_original_weights
+
+    state_dicts = []
+    sharded_checkpoints = sorted(sharded_checkpoints)
+    for sharded_checkpoint in sharded_checkpoints:
+        if not sharded_checkpoint.is_file():
+            continue
+        state_dicts.append(load_function(sharded_checkpoint.as_posix()))
+
+    parameters_metadata = metadata["parameters"]
+    transformation_specs_metadata = metadata["model_weight_transformation_specs"]
+
+    # We recreate the transformation specs from the metadata.
+    transformations_specs = []
+    for specs_metadata in transformation_specs_metadata:
+        specs = ModelWeightTransformationSpecs.from_metadata(specs_metadata)
+        transformations_specs.append(specs)
+
+    # We transform the sharded state dicts as follows:
+    # [state_dict_tp_rank_0, state_dict_tp_rank_1, ...]
+    #   ->  {
+    #           key: [state_dict_tp_rank_0[key], state_dict_tp_rank_1[key], ...],
+    #           for key in state_dict_tp_rank_0.keys()
+    #       }
+    paramater_names = state_dicts[0].keys()
+    sharded_state_dicts = {name: [state_dict[name] for state_dict in state_dicts] for name in paramater_names}
+
+    consolidated_state_dict = to_original_weights(transformations_specs, sharded_state_dicts, parameters_metadata)
+
+    return consolidated_state_dict
+
+
 @requires_neuronx_distributed
 def consolidate_model_parallel_checkpoints(checkpoint_dir: Path) -> Dict[str, "torch.Tensor"]:
     model_checkpoint_dir = checkpoint_dir / "model"
@@ -221,26 +259,37 @@ def consolidate_model_parallel_checkpoints(checkpoint_dir: Path) -> Dict[str, "t
     # Case 2: If no file was found, maybe the checkpoint was saved without xser.
     if not sharded_checkpoints:
         sharded_checkpoints = list(model_checkpoint_dir.glob("dp_rank_*.pt"))
-        load_function = torch.load
+        load_function = partial(torch.load, weights_only=True)
 
     if not sharded_checkpoints:
         raise ValueError(f"Could not find any sharded checkpoint in {model_checkpoint_dir.as_posix()}")
 
     pp_size = max((int(checkpoint_path.stem[-2:]) for checkpoint_path in sharded_checkpoints)) + 1
     checkpoints_grouped_by_pp_ranks = [[] for _ in range(pp_size)]
     metadatas = []
+    is_old_metadata = False
     for pp_rank in range(pp_size):
         for checkpoint_path in sharded_checkpoints:
             checkpoint_name = checkpoint_path.stem
             if int(checkpoint_name[-2:]) == pp_rank:
                 checkpoints_grouped_by_pp_ranks[pp_rank].append(checkpoint_path)
-        metadatas.append(torch.load(checkpoint_dir / f"mp_metadata_pp_rank_{pp_rank}.pt"))
+        if (checkpoint_dir / f"mp_metadata_pp_rank_{pp_rank}.pt").is_file():
+            is_old_metadata = True
+            metadatas.append(torch.load(checkpoint_dir / f"mp_metadata_pp_rank_{pp_rank}.pt"))
+        else:
+            with open(checkpoint_dir / f"mp_metadata_pp_rank_{pp_rank}.json") as fp:
+                metadatas.append(json.load(fp))
 
     consolidated_state_dict = {}
     for pp_rank, checkpoint_group_for_pp_rank in enumerate(checkpoints_grouped_by_pp_ranks):
-        consolidated_for_pp_rank = consolidate_tensor_parallel_checkpoints(
-            checkpoint_group_for_pp_rank, load_function, metadatas[pp_rank]
-        )
+        if is_old_metadata:
+            consolidated_for_pp_rank = old_consolidate_tensor_parallel_checkpoints(
+                checkpoint_group_for_pp_rank, load_function, metadatas[pp_rank]
+            )
+        else:
+            consolidated_for_pp_rank = consolidate_tensor_parallel_checkpoints(
+                checkpoint_group_for_pp_rank, load_function, metadatas[pp_rank]
+            )
         consolidated_state_dict.update(**consolidated_for_pp_rank)
 
     for key, tensor in consolidated_state_dict.items():