fix: Diverse model seeding across PP ranks by rrutmann · Pull Request #426 · Modalities/modalities

rrutmann · 2025-12-10T09:49:10Z

What does this PR do?

This PR gives a unique model seed for each pp rank, such that parameters are initialized differently across ranks.

General Changes

On each rank, add the pp rank to the model seed.

Breaking Changes

None

Checklist before submitting final PR

My PR is minimal and addresses one issue in isolation
I have merged the latest version of the target branch into this feature branch
I have reviewed my own code w.r.t. correct implementation, missing type hints, proper documentation, etc.
I have run a sample config for model training
I have checked that all tests run through (python tests/tests.py)
I have updated the internal changelog (CHANGELOG_DEV.md)

BlueCrescent

Overall LGTM.
Should we also explicitly allow seeding for the "model_initialized" component?
It will probably inherit the random state from the model_raw component but it seems a bit risky to me to assume that (also in the future) no other interaction with the random state happens between these two components (though, probably, only interactions that are asymmetrical between the ranks would be problematic). In particular, since we cannot guarantee the order in which the components are build, something like a dataloader component might even re-seed the random state.

le1nux

I checked the seeding (not the test) and from my understanding the changes do not provide the expected results (also what @BlueCrescent was hinting towards).

When we seed the raw model, the model weights are indeed deterministic at instantiation time. However, we also have model weight initialization which runs afterwards and would just override the weights / seeding.

Additionally, passing device_mesh to the model is coupling two components that should normally not know anything about each other.

I think we have to integrate the seeding to the weight initializer component and can also pass in the device_mesh there.

rrutmann · 2025-12-19T15:36:15Z

I checked the seeding (not the test) and from my understanding the changes do not provide the expected results (also what @BlueCrescent was hinting towards).

When we seed the raw model, the model weights are indeed deterministic at instantiation time. However, we also have model weight initialization which runs afterwards and would just override the weights / seeding.

Additionally, passing device_mesh to the model is coupling two components that should normally not know anything about each other.

I think we have to integrate the seeding to the weight initializer component and can also pass in the device_mesh there.

Yes that makes sense. I moved the seeding to the model initialization component

rrutmann · 2025-12-19T15:36:35Z

Overall LGTM. Should we also explicitly allow seeding for the "model_initialized" component? It will probably inherit the random state from the model_raw component but it seems a bit risky to me to assume that (also in the future) no other interaction with the random state happens between these two components (though, probably, only interactions that are asymmetrical between the ranks would be problematic). In particular, since we cannot guarantee the order in which the components are build, something like a dataloader component might even re-seed the random state.

See #426 (comment)

BlueCrescent

LGTM

le1nux

Generally good state.
Left a couple of comments.
My main concern is the global setting of the seed. A generator object might be favorable.

le1nux · 2025-12-19T17:21:44Z

    """NNModel class to define a base model."""

-    def __init__(self, seed: int = None, weight_decay_groups: Optional[WeightDecayGroups] = None):
+    def __init__(self, seed: Optional[int] = None, weight_decay_groups: Optional[WeightDecayGroups] = None):


Suggested change

def __init__(self, seed: Optional[int] = None, weight_decay_groups: Optional[WeightDecayGroups] = None):

def __init__(self, seed: int | None = None, weight_decay_groups: Optional[WeightDecayGroups] = None):

Do we even want to allow setting the seed here?
Could torch.manual_seed below have side effects with the new weight init implementation?

le1nux · 2025-12-19T17:22:46Z

    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: Optional[int] = None


Suggested change

seed: Optional[int] = None

seed: int | None = None

also other instances and docstring annotations should be adapted accordingly

le1nux · 2025-12-19T17:28:16Z

        return initialization
+
+    @staticmethod
+    def _set_seed(seed: Optional[int]):


this sets the seed globally. I think an even more robust way would be to use a local torch rng object.
Could this be integrated?

Something like:

g = torch.Generator() g.manual_seed(1234)

Otherwise, we might get into side-effects later on

le1nux · 2026-05-05T09:10:58Z

+        if seed is not None and has_parallelism_method(
+            device_mesh=device_mesh, parallelism_method=ParallelismDegrees.PP
+        ):
+            seed += get_parallel_rank(device_mesh=device_mesh, parallelism_method=ParallelismDegrees.PP)


This also means that depending on the parallelization method and also the number of parallelism degrees we get differently initialized layers even if the seed is the same.
Example:

DP with seed = 1, will have a differently initilized model than DP+PP with seed = 1.

One way to fix this is to always use the same seed but each PP stage has to skip the number of random values of the pervious stages.
However, I think this would be overkill and I would just place a warning when initialising the weights and parallelization methods are other than FSDP.
What do you think?

We could also give the hint that for full reproducibility a Distributed Checkpoint with FSDP directly after weight init.
Maybe we could even have an entry point for that in main.

something like:
modalities create_init_cp model_config.yaml

For some unit tests, this functionality would be nice to have anyways I think.

Any thoughts?

le1nux · 2026-05-05T09:14:19Z

+        std: float | str,
+        parameter_name_regexes: list[str],
+        hidden_dim: Optional[int] = None,
+        seed: Optional[int] = None,


Suggested change

seed: Optional[int] = None,

seed: int | None = None,

also various cases below

rrutmann added 5 commits December 8, 2025 14:18

fix: Initialize different weights across TP ranks

5f9f50e

feat: Consider pp rank for model seed

8c8c5ab

fix: Only consider PP rank for seeding

ab3daa0

test: Add test for different parameters on tp/pp ranks

62a1743

test: Check for equal parameters across data parallel processes

00a595b

BlueCrescent reviewed Dec 15, 2025

View reviewed changes

Comment thread tests/fsdp2_parallelization/test_parallel_seed_initialization.py Outdated

rrutmann requested a review from le1nux December 19, 2025 10:25

rrutmann self-assigned this Dec 19, 2025

le1nux requested changes Dec 19, 2025

View reviewed changes

rrutmann added 5 commits December 19, 2025 13:26

feat: Integrate seeding to model initialization

bf06da7

refactor: Move seeding logic to model initialization component

b137701

chore: Add seed and device_mesh to ComposedModelInitializationConfig

bff99f3

test: Adapt test to latest changes

98ff9db

chore: Remove old code

2e248ed

BlueCrescent approved these changes Dec 23, 2025

View reviewed changes

chore: Merge branch 'main' into seed

093fa33

le1nux requested changes May 5, 2026

View reviewed changes

	def __init__(self, seed: Optional[int] = None, weight_decay_groups: Optional[WeightDecayGroups] = None):
	def __init__(self, seed: int \| None = None, weight_decay_groups: Optional[WeightDecayGroups] = None):

Conversation

rrutmann commented Dec 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

General Changes

Breaking Changes

Checklist before submitting final PR

Uh oh!

BlueCrescent left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

le1nux left a comment

Choose a reason for hiding this comment

Uh oh!

rrutmann commented Dec 19, 2025

Uh oh!

rrutmann commented Dec 19, 2025

Uh oh!

BlueCrescent left a comment

Choose a reason for hiding this comment

Uh oh!

le1nux left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

rrutmann commented Dec 10, 2025 •

edited

Loading