SIGN IN SIGN UP

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