Refactor model configuration and parallelization ownership by tianyu-l · Pull Request #4810 · pytorch/torchtitan

tianyu-l · 2026-09-21T00:16:09Z

Human's note

remove ModelSpec and refactor fields into Model or Model.Config
remove per-model parallelize.py (including GraphTrainer), and instead use Model / Decoder / MultimodalModel / FluxModel to hold the logic, shared with model subclasses
- introduced Model._apply_fsdp() for the above reason
remove train_context as it only calls into get_spmd_context
make Module.parallelize() -> Module._parallelize() as it's not supposed to be overridden

Summary

remove ModelSpec and make each model class own its config, checkpoint adapter, optimizer hooks, and high-level parallelization behavior
consolidate common decoder and multimodal parallelization in BaseModel, Decoder, and MultimodalModel, while keeping Module._parallelize() as the low-level primitive
consolidate the regular GraphTrainer model lifecycle and remove its four per-model parallelize.py files
centralize SPMD context setup and make vocab-parallel log-prob computation take its TP process group explicitly
define max_context_length on decoder configs and set it in concrete model presets

Why

ModelSpec duplicated model identity outside the model class and allowed invalid combinations such as one model with another model checkpoint adapter or optimizer hooks. Per-model parallelization files also repeated behavior that belongs to shared model families. This change makes those relationships structural: selecting a model class determines its configuration and associated behavior, while model subclasses retain focused override points.

Backward compatibility is intentionally not preserved for the removed ModelSpec APIs.

Testing

pre-commit run --all-files
pytest tests/unit_tests/cpu -x -q (1169 passed, 2 skipped)
affected CPU tests (64 passed)
focused RL tests (28 passed)
affected GPU model tests (28 passed, 1 skipped)
distributed FSDP/sharding tests (9 passed)
GraphTrainer distributed regression tests (3 passed)
GraphTrainer test collection (499 tests)
Flux fake-PG FSDP+compile integration (passed)

The branch is rebased onto current main. This is a regular PR so the full upstream CI suite will run.

tianyu-l · 2026-09-21T01:39:43Z

            f"{type(self).__name__} must implement preprocess_inputs()."
        )

-    def verify_module_protocol(self) -> None:


capability merged into _parallelize() as it's executed recursively for every module already

tianyu-l · 2026-09-21T06:39:52Z

GraphTrainer is the same on main
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not extract specialized integer from data-dependent expression u41 (unhinted: u41).
cc @aditvenk

jinsooihm · 2026-09-21T18:25:48Z

 )
 from torchtitan.models.common.param_init import depth_scaled_std
-from torchtitan.models.deepseek_v3.parallelize import (
-    parallelize_deepseekv3 as parallelize_deepseek_v4,


It seems like removing this got rid of apply_fsdp_to_mtp_decoder used in parallelize_deepseekv3. And DSv4 just inherits from Decoder not MTPDecoder, so the specific fsdp behavior is lost.

Good catch. DeepSeek V4 was falling back to Decoder._apply_fsdp(), which does not treat its separate mtp_layers as decoder-layer FSDP units. I added a DeepSeekV4Model._apply_fsdp() override that follows MTPDecoder and calls apply_fsdp_to_mtp_decoder(), plus a regression test. Updated in 4c7729c.

wwwjn

This refactor looks great to me

wwwjn · 2026-09-21T19:17:12Z

+
+    multimodal_encoder_fqns: tuple[str, ...] = ()
+
+    def parallelize(


nit: in human note part, you note that this fuction should be named as _parallelize() ?

Model.parallelize() still exist as public entry point
Module.parallelize() becomes Module._parallelize()

…4829) ## Summary #4810 deleted `ModelSpec` and per-model `parallelize.py`. `Trainer.Config.model` (and `Controller.Config.model`) is the Suppress-annotated field. Four docs still used the old names, so copy-paste snippets fail with `TypeError: unexpected keyword argument 'model_spec'` or point at a file that no longer exists. Same class of stale-path cleanup as #4805. - `torchtitan/config/README.md`: `Trainer.Config(model=...)` and `Trainer.Config.model` in the CLI-freeze paragraph - `torchtitan/rl/README.md`: `Controller.Config(model=...)` - `torchtitan/overrides/README.md`: FQN example `model.layers.0.feed_forward` - `torchtitan/models/llama3/README.md`: point at `sharding.py` and `BaseModel.parallelize()` instead of `parallelize.py` Docs only. No behavior change. Fixes #4827. ## Test plan - [ ] Confirmed each replacement path exists on this branch (`Trainer.Config.model`, `Controller.Config.model`, `llama3/sharding.py`, `BaseModel.parallelize`) - [ ] Grepped the four files for leftover `model_spec=` Config keywords and `parallelize.py`

tianyu-l requested review from IvanKobzarev, SherlockNoMad, aditvenk, fegin, felipemello1, pzhan9, sanketpurandare, wconstab and wwwjn as code owners September 21, 2026 00:16

pytorch-bot Bot added ciflow/fake-pg Run 1-GPU Fake PG integration tests ciflow/rl labels Sep 21, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Sep 21, 2026

tianyu-l force-pushed the model_spec branch 2 times, most recently from fe8f3c1 to 4cf26f9 Compare September 21, 2026 01:25

tianyu-l commented Sep 21, 2026

View reviewed changes

tianyu-l requested review from acisseJZhong, anijain2305, jinsooihm and shuhuayu September 21, 2026 01:48

tianyu-l force-pushed the model_spec branch 2 times, most recently from cec4607 to 10d0fe9 Compare September 21, 2026 02:25

jinsooihm reviewed Sep 21, 2026

View reviewed changes

wwwjn approved these changes Sep 21, 2026

View reviewed changes

aditvenk approved these changes Sep 21, 2026

View reviewed changes

tianyu-l force-pushed the model_spec branch from 10d0fe9 to 4c7729c Compare September 21, 2026 22:57

Refactor model configuration and parallelization ownership

22b7a4c

tianyu-l force-pushed the model_spec branch from 4c7729c to 22b7a4c Compare September 21, 2026 23:08

tianyu-l merged commit b9e1d1f into main Sep 22, 2026
27 of 30 checks passed

YeonwooSung mentioned this pull request Sep 22, 2026

[docs] Stale ModelSpec / parallelize.py names after #4810 #4827

Closed

YeonwooSung mentioned this pull request Sep 22, 2026

[docs] Replace leftover ModelSpec / parallelize.py names after #4810 #4829

Merged

2 tasks

QIU023 mentioned this pull request Sep 23, 2026

[Kimi K3] Pipeline Parallelism support (text LLM side): PP Virtual Stage Cache Adapter for Attention Residual #4312

Open

Navigation Menu

Refactor model configuration and parallelization ownership - #4810

Refactor model configuration and parallelization ownership#4810
tianyu-l merged 1 commit into
mainfrom
model_spec

tianyu-l commented Sep 21, 2026 •

edited

Loading

Uh oh!

tianyu-l Sep 21, 2026

Uh oh!

tianyu-l commented Sep 21, 2026

Uh oh!

jinsooihm Sep 21, 2026

Uh oh!

tianyu-l Sep 21, 2026

Uh oh!

wwwjn left a comment

Uh oh!

wwwjn Sep 21, 2026

Uh oh!

tianyu-l Sep 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants


		multimodal_encoder_fqns: tuple[str, ...] = ()

		def parallelize(

Conversation

tianyu-l commented Sep 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Human's note

Summary

Why

Testing

Uh oh!

tianyu-l Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

tianyu-l commented Sep 21, 2026

Uh oh!

jinsooihm Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

tianyu-l Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

wwwjn left a comment

Choose a reason for hiding this comment

Uh oh!

wwwjn Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

tianyu-l Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

tianyu-l commented Sep 21, 2026 •

edited

Loading