[OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank by shengliangxu · Pull Request #2742 · NVIDIA/Model-Optimizer

Load the dense part sharded, and size the load warmup to each rank

shengliangxu

committed

72b7fd0

Import the dense-sharding helpers' modules at the top of the file

shengliangxu

committed

4058038

Keep the dense-sharding helpers' docs to what a maintainer needs

shengliangxu

committed

9c5b36b

Fail every rank when one fails preparing the meta-DCP load

shengliangxu

committed

8780ffd

Load the EP dense shards on the mesh's device, and say so briefly

shengliangxu

committed

3f1e23b

Load both aliases of a tied param in the EP dense load

shengliangxu

committed

c4ad0c3

Import safetensors at the top of the dense-sharding tests

shengliangxu

committed

b73b91b

Check for missing gradients before the EP test's collective comparisons

shengliangxu

committed

3d92a08

Hand the meta-DCP load's prepared state over by name

shengliangxu

committed

fa4ec4f

Import the layer-wildcard helper from either transformers location

shengliangxu

committed

ccd8b5b

Refuse a sharded dense part that from_pretrained did not leave on meta

shengliangxu

committed

e375bb1

Warn when the EP-aware warmup sizing falls back to transformers' own

shengliangxu

committed

e9b1c9b

Say why _fsdp2_wrap_modules imports where it does

shengliangxu

committed

0a67b30

Stand in a list for the unplaced keys in the dense-meta refusal test

shengliangxu

committed

64c5121

Import the dense-wrap helpers at module scope

shengliangxu

committed

f51aedf

Warn when the warmup sizing hook has other parameters, and name EP_DEGREE=1

shengliangxu

committed

b579ddd

Say which VLMs keep their vision blocks unwrapped, and which buffers load as usual

shengliangxu

committed

a79a929

Move the dense-sharding and warmup tests into test_model_load_utils.py

shengliangxu

committed

9fe4dea

Materialize the sharded dense part on the GPU each rank is bound to

shengliangxu

committed

8a21375

Fail every rank when one rank fails to wrap the sharded dense part

shengliangxu

committed

855600d

Say so when transformers has no warmup sizing hook to replace

shengliangxu

committed

dd04431

Navigation Menu

Repository navigation

[OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank - #2742

[OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank#2742
shengliangxu wants to merge 21 commits into
shengliangx/hfep-6a2-ep-loader-corefrom
shengliangx/hfep-6b-ep-dense-sharding

Commits on Nov 1, 2026

Load the dense part sharded, and size the load warmup to each rank

Import the dense-sharding helpers' modules at the top of the file

Keep the dense-sharding helpers' docs to what a maintainer needs

Fail every rank when one fails preparing the meta-DCP load

Load the EP dense shards on the mesh's device, and say so briefly

Load both aliases of a tied param in the EP dense load

Import safetensors at the top of the dense-sharding tests

Check for missing gradients before the EP test's collective comparisons

Hand the meta-DCP load's prepared state over by name

Import the layer-wildcard helper from either transformers location

Refuse a sharded dense part that from_pretrained did not leave on meta

Warn when the EP-aware warmup sizing falls back to transformers' own

Say why _fsdp2_wrap_modules imports where it does

Stand in a list for the unplaced keys in the dense-meta refusal test

Import the dense-wrap helpers at module scope

Warn when the warmup sizing hook has other parameters, and name EP_DEGREE=1

Say which VLMs keep their vision blocks unwrapped, and which buffers load as usual

Move the dense-sharding and warmup tests into test_model_load_utils.py

Materialize the sharded dense part on the GPU each rank is bound to

Fail every rank when one rank fails to wrap the sharded dense part

Say so when transformers has no warmup sizing hook to replace