Skip to content

feat(fsdp2): load model on meta device for non-rank0 ranks (0 CPU RAM per worker) #9982

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cben484 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat(fsdp2): load model on meta device for non-rank0 ranks (0 CPU RAM per worker) #9982

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
60b1dcc
feat(fsdp2): load model on meta device for non-rank0 ranks (0 CPU RAM)
cben484 Aug 24, 2026
3e7771b
style: apply yapf formatting
cben484 Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions swift/model/register.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,22 @@
transformers_5 = version.parse(transformers.__version__) >= version.parse('5.0.0.dev')


def _fsdp2_use_meta_loading() -> bool:
"""Whether non-rank0 ranks should load the model on the meta device under FSDP2 (0 CPU RAM).

The weights are later synced from rank0 by accelerate's ``cpu_ram_efficient_loading``
path during ``prepare``. Can be disabled with env ``SWIFT_FSDP2_META_LOADING=0``.
"""
if os.environ.get('SWIFT_FSDP2_META_LOADING', '1') == '0':
return False
if os.environ.get('FSDP_VERSION') != '2': # set by SftArguments._init_fsdp
return False
try:
return int(os.environ.get('WORLD_SIZE', '1')) > 1
except ValueError:
return False


def register_model(model_meta: ModelMeta, *, exist_ok: bool = False) -> None:
"""
model_type: The unique ID for the model type. Models with the same model_type share
Expand Down Expand Up @@ -633,6 +649,21 @@ def get_model_processor(
new_special_tokens=new_special_tokens,
model_kwargs=model_kwargs,
**kwargs)
if load_model and _fsdp2_use_meta_loading() and int(os.environ.get('RANK', '0')) != 0:
from accelerate import init_empty_weights
logger.info('FSDP2 non-rank0: loading the model on the meta device (0 CPU RAM), '
'weights will be synced from rank0 by accelerate during prepare.')
# Mask ACCELERATE_USE_FSDP during loading: transformers' FSDP non-rank0 branch
# (modeling_utils._move_missing_keys) would otherwise materialize meta params
# as CPU zeros (full model in host RAM per rank). device_map was already
# resolved earlier in this function, so masking here is safe.
use_fsdp = os.environ.pop('ACCELERATE_USE_FSDP', None)
try:
with init_empty_weights():
return loader.load()
finally:
if use_fsdp is not None:
os.environ['ACCELERATE_USE_FSDP'] = use_fsdp
return loader.load()


Expand Down