Pipeline parallel naive inference (#47289)
* add fsdp plan to 2 models for now * add tests fsdp mixin * linting * refactor test fsdp mixin * test fsdp mixin cleaning * remove fsdp policy in tests + trim down further * test fsdp clean * restore test_modeling_utils * linting * start trim down stuff * fix * breaking: cleaning modeling_utils.py * load path with fsdp (dtensor) and tp (old tp) is linked * linting * add saving * styling * fix tp ci * add fsdp to ci * linting * pick one model only for this PR * restore * trigger fsdp ci * doc cleaning + tp_size remove * fix tp ci for ep * edit doc * move distributed function to utils + guarding * linting * expand_fsdp_plan iterate over modules * comment about tie embedding * add comment tied embedding * add DistributedMixin * some cleaning * cleaning + comment * rename function for clarity * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * doc * comment * linting * refactor * abstract to mixin * typo * Add FSDP plans to all models from distributed branch. Port base_model_fsdp_plan and ForCausalLM _fsdp_plan entries from PR #46269 and expand FSDP distributed test coverage to the pilot model subset. * fsdp plans * linting * linting * 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. * Wire DistributedConfig through from_pretrained and save_pretrained (FSDP orchestration 2/3). Route distributed loading and saving through DistributedMixin, migrate TP tests and docs off tp_plan="auto", and add FSDP gather/DCP save paths. * Add FSDP CI and end-to-end FSDP tests (FSDP orchestration 3/3). Add FSDPTesterMixin, cohere2_moe base_fsdp_plan, dedicated fsdp_ci job, and pytest markers for distributed FSDP load/save/generation coverage. * 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 * revert * revert * post_init() parallel plan move to mixin * revert tp mixin * add save/load * adapt distributed config to handle pp_size * edit validate * init pipeline parallelism * avoid loading weights with TP when pp_size>1 * cleaning args in initialize_pipeline_parallelism * add very dumb layers splitting * add test pipeline parallel * draft pp + generate test to be clean * add weight loading logging * cleaning * remove comm_device to make it an attribute * begin to refactor * avoid loading report when tp is enabled * cleaning pp * better cleaning pp * cleaning pp * better cleaning broadcast_from_last * cleaning pp * pp tests no mixing * handle tied embeddings * linting * doc * unecessary if * comment * linting * make sure pp doesnt go into tp shard and distribute * only FSDP save/load for now * revert * refactor * modular * ea * begin migration TP * clean up * add it to pretrained model * linting * fix the test by moving in init class the fsdp plan instead of post init * edit * replace everything * adapt distributed config to handle pp_size * init pipeline parallelism * avoid loading weights with TP when pp_size>1 * cleaning args in initialize_pipeline_parallelism * add very dumb layers splitting * add test pipeline parallel * draft pp + generate test to be clean * add weight loading logging * cleaning * remove comm_device to make it an attribute * begin to refactor * avoid loading report when tp is enabled * cleaning pp * better cleaning pp * cleaning pp * better cleaning broadcast_from_last * cleaning pp * pp tests no mixing * handle tied embeddings * linting * doc * unecessary if * comment * linting * make sure pp doesnt go into tp shard and distribute * begin migration TP * clean up * revert merge conflict * local params for forward * cleaning * migration from integration.tensor_parallel to distributed.tensor_parallel * revert models * revert * fix * breaking: cleaner way to_local for quantize weight(almost done need to fix backward) * fix all tests * TP dtensor handle natively the to_local() for deepgem + fp8 (#47634) * Refactor distributed tensor_parallel module and TP mixin tests * better approach to to_local * pass as class method * cleaning PackedColwise * remove the use of to_local in kernels to offload this task to TensorParallel * rowwise bias after redistribute * make it easier to understand * requires grad only for floating point numbers * fix tp meagamoe fp8 with dtensor * linting * revert gitignore * fix run slow quantization ci * fix deepseek v4 ep backward tests * fix deepseek glm4 moe tp backward pass * Refactor tensor parallel loading logic by removing unused tp_plan handling and updating is_dtensor type hint for better type safety. * linting * bench: dtensor vs legacy TP * bench: dtensor vs legacy TP (#47728) * bench: dtensor vs legacy TP * clean rowise * refactor colwise * comment * comment * refactor MoEExperParallel * ruff * add rowwise input * remove helper script * remove dead code in mxfp4 * remove old TP * cleaning * cleaning * claening * add test_shard_tensor_shape_consistency * cleaning * small fix * typo * comment * clean * typo * rename colwise_gather_output to colwise_rep * keep tp import to avoid breaking chnges * remove async_op=True for redistribute * linting * renaming MoeExpertsParallel class * ep router doc * mlinter * add todo * fix * new name * guard import * cache inspect signature * merge init and from_device_mesh * add todo * avoid nesting * pp stages logs only ad INFO level * better doc func * linting * linting * fix --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
F
Ferdinand Mom committed
b0c8d427be61b868b98d86435c529243b2da75db
Parent: da7234a
Committed by GitHub <noreply@github.com>
on 8/25/2026, 2:03:57 PM