Add distributed runtime utils and DistributedMixin - #47352
Conversation
…n 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points.
3outeille
changed the title
3outeille
changed the title
|
run-slow: deepseek_v4, glm_moe_dsa, gpt_oss |
|
This comment contains models: ["models/deepseek_v4", "models/glm_moe_dsa", "models/gpt_oss"] |
CI ResultsCommit Info
The test failure analysis could not be completed. Please check the workflow run for details. |
| tp_plan="auto", | ||
| distributed_config=DistributedConfig(enable_expert_parallel=True), | ||
| distributed_config=DistributedConfig( | ||
| tp_size=int(os.environ["WORLD_SIZE"]), |
There was a problem hiding this comment.
tp size could be auto no?
There was a problem hiding this comment.
we need to specify the tp_size because having 8 gpus doesnt mean we want to automatically split on 8 gpus. However tp_plan is automatically triggered if tp_size > 1
| if not is_torch_available(): | ||
| raise RuntimeError("PyTorch is required to use DistributedConfig.") | ||
|
|
||
| if not torch.distributed.is_available() or not torch.distributed.is_initialized(): | ||
| raise RuntimeError( | ||
| "torch.distributed must be initialized before using DistributedConfig with tp_size > 1 or " | ||
| "fsdp_size > 1. Call dist.init_process_group(...) first, or launch with torchrun." | ||
| ) |
There was a problem hiding this comment.
let's make sure we don't already gate this somewhere else
There was a problem hiding this comment.
i'll defer to another PR as for now TP + FSDP is not linked together. They are both initialising the process group which is redudant. Once it will be unified, i'll remove gate once and for all
|
[For maintainers] Suggested jobs to run (before merge) run-slow: deepseek_v4, glm_moe_dsa, gpt_oss |
CI recapDashboard: View test results in Grafana |
@3outeille Isn't this a backwards incompatible change? I thought the idea with Transformers v5 was to minimize those. This change doesn't even come with a deprecation cycle or useful error message for users to update their code. |
* Add distributed runtime utils and DistributedMixin (FSDP orchestration 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet. * addd ep_plan * restore validate module * Wire DistributedConfig through from_pretrained and save_pretrained. Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points. * revert * inline distribute_model * revert * remove saving/loading * leaner mixin * downgrade torch version guarding * remove * linting * post_init() parallel plan move to mixin * revert tp mixin
|
@BenjaminBossan fair point indeed, this is backwards incompatible and I should have shipped a deprecation path with it. I would be more mindful next time as the distributed features are not stable yet |
* Add distributed runtime utils and DistributedMixin (FSDP orchestration 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet. * addd ep_plan * restore validate module * Wire DistributedConfig through from_pretrained and save_pretrained. Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points. * revert * inline distribute_model * revert * remove saving/loading * leaner mixin * downgrade torch version guarding * remove * linting * post_init() parallel plan move to mixin * revert tp mixin
4 participants
Summary
Part of FSDP orchestration stack (1/2). Replaces #46990.
distributed/utils.py(mesh init,distribute_model)DistributedMixinwith plan properties and orchestration hooksDistributedConfigwith deferredvalidate()and FSDP+TP guardapply_fully_sharded_data_parallel→apply_fully_sharded_data_parallelism(torch>=2.7)PreTrainedModelto inheritDistributedMixin(plan properties moved out)distribute_model→apply_tensor_parallelismin tensor_parallel.pyprepare_distribute_model/maybe_distribute_modelintofrom_pretrainedtp_plan/tp_sizekwargs fromfrom_pretrained; usedistributed_config=DistributedConfig(tp_size=N)insteadDistributedConfig