fix(fsdp): set ACCELERATE_USE_FSDP so device_map and cpu_ram_efficient_loading take effect

[cben484]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Problem

When training with --fsdp fsdp2, every rank loads the full model weights onto its own device, causing two silent failures:

1. Immediate OOM at load time for large models: when the weights exceed a single card's memory (e.g., 72B bf16 ≈ 144G, or a 35B MoE ≈ 67G vs 64G HBM), from_pretrained OOMs during loading and never reaches the FSDP2 sharding stage.
2. cpu_ram_efficient_loading silently has no effect: swift/config/fsdp2.json ships with "cpu_ram_efficient_loading": true, intended to make rank 0 load the weights while other ranks init on meta device. Because the switch is never wired up, all ranks each materialize a full copy of the weights in CPU RAM — N cards means N × full-model host memory (measured ~136G RSS per process on 16 cards, exhausting host memory and triggering the global OOM killer).

Small models (e.g., the 4B model in the official GRPO example) don't OOM, which is why this has gone unnoticed so far.

Root Cause

The ACCELERATE_USE_FSDP environment variable is read in two places:

• swift's own get_default_device_map() (swift/model/utils.py):

def get_default_device_map():
    if is_deepspeed_zero3_enabled() or os.environ.get('ACCELERATE_USE_FSDP', 'False') == 'true':
        return None
    ...
    return 'auto' if is_mp() else f'cuda:{local_rank}'   # NPU: f'npu:{local_rank}'

• transformers' is_fsdp_enabled() (the gate for the FSDP rank0-only / meta-device init load path)

The variable is normally set by accelerate when the Accelerator is created, but swift loads the model before the Trainer/Accelerator is created, so the variable doesn't exist yet; _init_fsdp() only sets FSDP_VERSION. As a result, both readers conclude FSDP is disabled:

• get_default_device_map() returns cuda:{local_rank} / npu:{local_rank} → every rank loads the full weights
• transformers skips the rank0-only path → cpu_ram_efficient_loading has no effect

ZeRO-3 is unaffected because is_deepspeed_zero3_enabled() has its own dedicated branch — which is also why this bug only shows up on the FSDP path.

Fix

Set os.environ['ACCELERATE_USE_FSDP'] = 'true' in _init_fsdp():

• Scoped safely: _init_fsdp() returns early when --fsdp is unset, so non-FSDP users are unaffected
• Correct timing: runs during argument processing, before model loading
• One change, two beneficiaries: wires up both swift's get_default_device_map() (no more device_map) and transformers' is_fsdp_enabled() (rank0-only load path)

Why not patch get_default_device_map() itself (e.g., also check FSDP_VERSION)? That would only fix the device_map side; the transformers rank0-only path would still be dead. Setting the env var covers both readers at once.

Reproduce

Any model whose weights exceed one card's memory reproduces this (GPU/NPU alike):

NPROC_PER_NODE=8 swift sft \
    --model Qwen/Qwen2.5-72B-Instruct \
    --fsdp fsdp2 \
    --tuner_type lora \
    --dataset <any> --max_steps 2

Two observations:

1. The log shows model_kwargs: {'device_map': 'cuda:0', ...} (expected: None)
2. N progress bars (one per rank) of Loading weights: 100% full loading, followed by an OOM during loading

Small models are also affected (no OOM, but wrong behavior): N full-load progress bars + each rank's CPU RSS ≈ full model size.

Scope

• Only affects --fsdp users: device_map behavior fixed (no more full loading) + cpu_ram_efficient_loading path wired up
• Zero changes for non-FSDP paths (early-return branch)

Experiment results

Measured on Ascend 910B (64G HBM × 8), model Qwen/Qwen3.5-35B-A3B (MoE, bf16 ≈ 67G), swift rlhf DPO + LoRA + FSDP2:

	Before fix	After fix
model_kwargs	{'device_map': 'npu:0'}	{'device_map': None}
Weight loading	Every card loads full weights onto NPU; OOMs at 63% (59.9G / 61G usable)	Weights stay on CPU; only sharded onto NPU after FSDP2 fully_shard; no OOM
Training	Unreachable	8-card FSDP2 trains normally, 28.3 GiB/card, DPO/KTO smoke tests pass (loss & rewards metrics healthy)

Note on the CPU side (transformers 5.12.1): from_pretrained does still have an FSDP non-rank0 branch (modeling_utils.py, _move_missing_keys, gated by is_fsdp_enabled() and not is_local_dist_rank_0()), but it materializes every parameter with torch.zeros_like(param, device='cpu') — i.e. the full model in host RAM per rank (~136G peak incl. prepare-time copies for a 35B bf16 model), with real weights broadcast from rank0 by accelerate afterwards. It is not a zero-memory load, and it is not keyed on cpu_ram_efficient_loading (that identifier no longer appears in modeling_utils.py). So even with this PR, per-rank CPU usage stays at ~model-size: 8 cards × ~136G ≈ 1.1T is fine on a 2T host, 16 cards gets tight. A follow-up PR will make non-rank0 ranks load on the meta device (validated on 16× Ascend 910B: host RAM peak 187G total instead of ~2.2T, identical training metrics vs full loading), letting accelerate's cpu_ram_efficient_loading sync be the only materialization path.

Checklist

• Verified on Ascend NPU (35B MoE, FSDP2 8-card DPO/KTO pass)
• Zero impact on non-FSDP paths (_init_fsdp() early return)
• GPU regression (no GPU env at hand; reviewer verification of the repro steps would be appreciated)