Merge pull request #426 from Modalities/seed

fix: Diverse model seeding across PP ranks

Author: rrutmann

config_files/training/config_lorem_ipsum_long_fsdp2.yaml

@@ -13,6 +13,7 @@ settings:
     checkpoint_saving_path: data/checkpoints
     train_dataset_path: ./data/lorem_ipsum_long.pbin
     test_dataset_path: ./data/lorem_ipsum.pbin
+    experiments_root_path: ${modalities_env:experiments_root_path}
   intervals:
     training_log_interval_in_steps: 1
     checkpointing_interval_in_steps: 32
@@ -221,6 +222,7 @@ initialized_model:
         mean: 0.0
         std: 0.02
         num_layers: ${model_raw.config.n_layer}
+        multi_device_generator_policy: error
 
 fsdp_model:
   component_key: model

config_files/training/config_lorem_ipsum_long_fsdp2_pp_tp.yaml

@@ -223,6 +223,10 @@ initialized_model:
         mean: 0.0
         std: 0.02
         num_layers: ${model_raw.config.n_layer}
+        seed: 42
+        device_mesh:
+          instance_key: device_mesh
+          pass_type: BY_REFERENCE
 
 scheduled_pipeline:
   component_key: pipeline
@@ -315,7 +319,6 @@ model_raw:
   component_key: model
   variant_key: gpt2
   config:
-    seed: 42
     use_meta_device: true
     use_weight_tying: false
     sample_key: ${settings.referencing_keys.sample_key}

docs/components/components.md

@@ -17,6 +17,8 @@
 |---------------|--------------------|----------------|---------------|---------------------|-------------|
 | model_initialization | composed | [ComposedInitializationRoutines.get_composed_model_initializer](../../src/modalities/nn/model_initialization/composed_initialization.py)| [ComposedModelInitializationConfig](../../src/modalities/nn/model_initialization/composed_initialization.py) | [ModelInitializationIF](../../src/modalities/nn/model_initialization/initialization_if.py) | Component for initializing model weights in place |
 
+The composed initializer supports seeded weight initialization for reproducibility within a fixed topology. When pipeline parallelism is active, Modalities offsets the initialization seed by pipeline stage rank to avoid identical stage-local weights. As a result, the same seed can produce different initialized weights for different pipeline-parallel topologies. For topology-independent reproducibility, create and reuse a distributed checkpoint directly after weight initialization.
+
 ## Losses
 
 |Component type | Component Version  | Implementation | Configuration | Component Interface | Description |

pyproject.toml

@@ -124,6 +124,7 @@ line-length = 120
 
 [tool.pytest.ini_options]
 addopts = "--cov=src --cov-report term --cov-report html"
+#addopts = "-ra" # Enable this instead of line above for reliable VS Code test debugging (without coverage)
 
 [tool.coverage.run]
 branch = true

src/modalities/config/config.py

@@ -7,8 +7,8 @@
 from omegaconf import OmegaConf, Resolver
 from pydantic import BaseModel, ConfigDict, Field, FilePath, PositiveInt, field_validator, model_validator
 from torch.distributed.fsdp import ShardingStrategy
-from transformers import GPT2TokenizerFast
-from transformers.models.llama.tokenization_llama_fast import LlamaTokenizerFast
+from transformers import GPT2Tokenizer as GPT2TokenizerFast
+from transformers import LlamaTokenizer as LlamaTokenizerFast
 from typing_extensions import deprecated
 
 from modalities.config.lookup_enum import LookupEnum

src/modalities/conversion/gpt2/modeling_gpt2.py

@@ -40,7 +40,14 @@
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
-from transformers.utils.generic import check_model_inputs
+
+try:
+    from transformers.utils.generic import check_model_inputs
+except ImportError:
+
+    def check_model_inputs(func: Callable) -> Callable:
+        return func
+
 
 from modalities.conversion.gpt2.configuration_gpt2 import GPT2Config
 

src/modalities/models/gpt2/gpt2_model.py

@@ -342,7 +342,6 @@ class GPT2LLMConfig(BaseModel):
         ffn_norm_config (LayerNormWrapperConfig): Config for normalization of the feed-forward network.
         lm_head_norm_config (LayerNormWrapperConfig): Config for normalization of the language model head.
         use_weight_tying (bool): Whether to use weight tying.
-        seed: Optional[int] = None: The random seed for reproducibility.
         enforce_swiglu_hidden_dim_multiple_of (int): If specified, enforces the hidden dimension
             in the SwiGLU layer to be a multiple of this value. Note that this is only relevant if the
             activation_type is SwiGLU. Defaults to 256.
@@ -370,7 +369,6 @@ class GPT2LLMConfig(BaseModel):
     ffn_norm_config: LayerNormWrapperConfig
     lm_head_norm_config: LayerNormWrapperConfig
     use_weight_tying: bool
-    seed: Optional[int] = None
     enforce_swiglu_hidden_dim_multiple_of: int = 256
 
     @model_validator(mode="after")
@@ -837,7 +835,6 @@ def __init__(
         ffn_norm_config: LayerNormWrapperConfig,
         lm_head_norm_config: LayerNormWrapperConfig,
         use_weight_tying: bool,
-        seed: Optional[int] = None,
         enforce_swiglu_hidden_dim_multiple_of: int = 256,
     ):
         """
@@ -862,7 +859,6 @@ def __init__(
             attention_norm_config (LayerNormWrapperConfig): Config for the attention normalization module.
             ffn_norm_config (LayerNormWrapperConfig): Config for the feed-forward network normalization module.
             lm_head_norm_config (LayerNormWrapperConfig): Config for the language model head normalization module.
-            seed (int, optional): The random seed. Defaults to None.
             use_weight_tying (bool): Whether to use weight tying.
             enforce_swiglu_hidden_dim_multiple_of (int): Enforces
                 the hidden dimension in the SwiGLU layer to be a multiple of this value.
@@ -873,7 +869,7 @@ def __init__(
             "embedding": [".wte", ".wpe"],
             "layernorm": [".attention_norm", ".ffn_norm", ".lm_head_norm"],
         }
-        super().__init__(weight_decay_groups=weight_decay_groups, seed=seed)
+        super().__init__(weight_decay_groups=weight_decay_groups)
         self.sample_key = sample_key
         self.prediction_key = prediction_key
         self.sequence_length = sequence_length

src/modalities/models/model.py

@@ -26,16 +26,13 @@ class ActivationType(str, Enum):
 class NNModel(nn.Module):
     """NNModel class to define a base model."""
 
-    def __init__(self, seed: int = None, weight_decay_groups: Optional[WeightDecayGroups] = None):
+    def __init__(self, weight_decay_groups: Optional[WeightDecayGroups] = None):
         """
         Initializes an NNModel object.
 
         Args:
-            seed (int, optional): The seed value for random number generation. Defaults to None.
             weight_decay_groups (Optional[WeightDecayGroups], optional): The weight decay groups. Defaults to None.
         """
-        if seed is not None:
-            torch.manual_seed(seed)
         self._weight_decay_groups = weight_decay_groups if weight_decay_groups is not None else {}
         super(NNModel, self).__init__()
 

src/modalities/models/model_factory.py

@@ -615,7 +615,6 @@ def get_gpt2_model(
         lm_head_norm_config: LayerNormWrapperConfig,
         use_weight_tying: bool,
         use_meta_device: Optional[bool] = False,
-        seed: Optional[int] = None,
         enforce_swiglu_hidden_dim_multiple_of: int = 256,
     ) -> GPT2LLM:
         config = dict(
@@ -637,7 +636,6 @@ def get_gpt2_model(
             attention_norm_config=attention_norm_config,
             ffn_norm_config=ffn_norm_config,
             lm_head_norm_config=lm_head_norm_config,
-            seed=seed,
             use_weight_tying=use_weight_tying,
             enforce_swiglu_hidden_dim_multiple_of=enforce_swiglu_hidden_dim_multiple_of,
         )

src/modalities/nn/model_initialization/composed_initialization.py

@@ -1,17 +1,26 @@
 from typing import Optional
 
+import torch
 import torch.nn as nn
 from pydantic import BaseModel, ConfigDict, Field, model_validator
+from torch.distributed.device_mesh import DeviceMesh
 from typing_extensions import Annotated
 
-from modalities.config.pydantic_if_types import PydanticModelInitializationIFType
+from modalities.config.pydantic_if_types import PydanticDeviceMeshIFType, PydanticModelInitializationIFType
 from modalities.nn.model_initialization.initialization_if import ModelInitializationIF
-from modalities.nn.model_initialization.initialization_routines import InitializationRoutines
+from modalities.nn.model_initialization.initialization_routines import (
+    InitializationRoutines,
+    MultiDeviceGeneratorPolicy,
+)
 from modalities.nn.model_initialization.parameter_name_filters import (
     NAMED_PARAMETER_INIT_GROUPS,
     SupportWeightInitModels,
     WeightInitTypes,
 )
+from modalities.running_env.fsdp.device_mesh import ParallelismDegrees, get_parallel_rank, has_parallelism_method
+from modalities.utils.logger_utils import get_logger
+
+logger = get_logger(__name__)
 
 
 class ModelInitializerWrapperConfig(BaseModel):
@@ -30,6 +39,9 @@ class ComposedModelInitializationConfig(BaseModel):
     std: Annotated[float, Field(strict=True, ge=0.0)] | str  # can be float or "auto"
     hidden_dim: Optional[Annotated[int, Field(strict=True, gt=0)]] = None
     num_layers: Optional[Annotated[int, Field(strict=True, gt=0)]] = None
+    seed: int | None = None
+    multi_device_generator_policy: MultiDeviceGeneratorPolicy = MultiDeviceGeneratorPolicy.WARN
+    device_mesh: Optional[PydanticDeviceMeshIFType] = None
 
     # avoid warning about protected namespace 'model_', see
     # https://docs.pydantic.dev/2.7/api/config/#pydantic.config.ConfigDict.protected_namespaces
@@ -87,6 +99,24 @@ def initialize_in_place(self, model: nn.Module):
 
 
 class ComposedInitializationRoutines:
+    @staticmethod
+    def _warn_pp_topology_dependent_seed(device_mesh: Optional[DeviceMesh], seed: Optional[int]) -> None:
+        if seed is None or not has_parallelism_method(
+            device_mesh=device_mesh, parallelism_method=ParallelismDegrees.PP
+        ):
+            return
+
+        if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0:
+            return
+
+        logger.warning(
+            "Seeded weight initialization is topology-dependent when pipeline parallelism is active. "
+            "Modalities offsets the initialization seed by PP rank to avoid identical stage-local weights, "
+            "so the same seed can produce different initialized weights for different PP configurations. "
+            "For topology-independent reproducibility, create and reuse a distributed checkpoint directly "
+            "after weight initialization."
+        )
+
     @staticmethod
     def get_model_initializer_wrapper(model_initializers: list[ModelInitializationIF]) -> ModelInitializationIF:
         initializer_wrapper = ModelInitializerWrapper(model_initializers)
@@ -98,8 +128,11 @@ def get_composed_model_initializer(
         weight_init_type: WeightInitTypes,
         mean: float,
         std: float | str,
-        hidden_dim: Optional[int] = None,
-        num_layers: int = None,
+        hidden_dim: int | None = None,
+        num_layers: int | None = None,
+        device_mesh: Optional[DeviceMesh] = None,
+        seed: int | None = None,
+        multi_device_generator_policy: MultiDeviceGeneratorPolicy = MultiDeviceGeneratorPolicy.WARN,
     ) -> ModelInitializationIF:
         """This initialization allows to intialize a model with plain, scaled or scaled_embed initialization.
         Note that plain initialization is always performed in the beginning. In case of scaled_embed,
@@ -114,36 +147,64 @@ def get_composed_model_initializer(
                 Defaults to None.
             num_layers (int, optional): Number of layers in the model (required for scaled and scaled_embed only).
                 Defaults to None.
+            device_mesh (Optional[DeviceMesh], optional): Device mesh used for parallelization.
+            seed (Optional[int], optional): Seed for random initialization. Defaults to None. When pipeline
+                parallelism is active, the effective seed is offset by PP rank to avoid identical stage-local
+                initialization, so the same seed does not guarantee identical initialized weights across different
+                PP topologies.
+            multi_device_generator_policy (MultiDeviceGeneratorPolicy, optional): Behavior when
+                initialization creates per-device RNG generators for more than one device in the same process.
+                Defaults to MultiDeviceGeneratorPolicy.WARN.
 
         Returns:
             ModelInitializationIF: The Weight Initializer performing the initialization as specified.
         """
+        ComposedInitializationRoutines._warn_pp_topology_dependent_seed(device_mesh=device_mesh, seed=seed)
+
+        # Set different random seed for each PP rank to ensure diversity
+        if seed is not None and has_parallelism_method(
+            device_mesh=device_mesh, parallelism_method=ParallelismDegrees.PP
+        ):
+            assert device_mesh is not None
+            seed += get_parallel_rank(device_mesh=device_mesh, parallelism_method=ParallelismDegrees.PP)
+
         model_initializers = []
 
         # plain
         plain_parameter_name_regexes = NAMED_PARAMETER_INIT_GROUPS[model_type][WeightInitTypes.PLAIN]
         plain_init = InitializationRoutines.get_plain_initialization(
-            mean=mean, std=std, hidden_dim=hidden_dim, parameter_name_regexes=plain_parameter_name_regexes
+            mean=mean,
+            std=std,
+            hidden_dim=hidden_dim,
+            parameter_name_regexes=plain_parameter_name_regexes,
+            seed=seed,
+            multi_device_generator_policy=multi_device_generator_policy,
         )
         working_std = plain_init.std
         model_initializers.append(plain_init)
 
         if weight_init_type in [WeightInitTypes.SCALED, WeightInitTypes.SCALED_EMBED]:
             # scaled
+            assert num_layers is not None
             scaled_parameter_name_regexes = NAMED_PARAMETER_INIT_GROUPS[model_type][WeightInitTypes.SCALED]
             scaled_init = InitializationRoutines.get_scaled_initialization(
                 mean=mean,
                 std=working_std,
                 num_layers=num_layers,
                 parameter_name_regexes=scaled_parameter_name_regexes,
+                seed=seed,
+                multi_device_generator_policy=multi_device_generator_policy,
             )
             model_initializers.append(scaled_init)
 
         if weight_init_type == WeightInitTypes.SCALED_EMBED:
             # scaled embed
             scaled_embed_parameter_name_regexes = NAMED_PARAMETER_INIT_GROUPS[model_type][WeightInitTypes.SCALED_EMBED]
             scaled_embed_init = InitializationRoutines.get_scaled_embed_initialization(
-                mean=mean, parameter_name_regexes=scaled_embed_parameter_name_regexes
+                mean=mean,
+                parameter_name_regexes=scaled_embed_parameter_name_regexes,
+                seed=seed,
+                multi_device_generator_policy=multi_device_generator_policy,
             )
             model_initializers.append(scaled_embed_init)