feat(fsdp): add --fsdp-reshard-after-forward for ZeRO-2 style training

[Rockdu]

What

• Add --fsdp-reshard-after-forward / --no-fsdp-reshard-after-forward (default: on, i.e. current behavior).
• apply_fsdp2 now forwards the value to every fully_shard call (per-block wraps and the root) as reshard_after_forward.

Why

FSDP2 defaults to resharding parameters after forward (ZeRO-3), which re-all-gathers every wrapped block during backward. When the unsharded parameters fit in memory, --no-fsdp-reshard-after-forward keeps them gathered between forward and backward (ZeRO-2), removing the backward all-gather at the cost of holding unsharded parameters through the step.

The flag deliberately mirrors fully_shard's own reshard_after_forward argument name (prefixed with fsdp-) instead of introducing a zero2 alias, so it maps 1:1 to the PyTorch API.

Files

• miles/utils/arguments.py — new --fsdp-reshard-after-forward tri-state flag (BooleanOptionalAction, default True).
• miles/backends/fsdp_utils/actor.py — pass it via fsdp_kwargs to all fully_shard calls; include it in the FSDP wrap log line.
• tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_integration_worker.py — add the new field to the bare Namespace this worker feeds apply_fsdp2.
• docs/user-guide/cli-reference.md — document the flag in the training-backend table.

Checklist

• pre-commit run passes (run on the touched files; --all-files not run)
• Added/updated tests for new behaviour — no dedicated test for the ZeRO-2 path (would need a multi-GPU run); only the existing integration worker's args were updated
• pytest -x is green — not run locally (this environment lacks sglang, which the test imports require)
• python3 train.py --help still parses — partially verified: the new argument's parse behavior (default / --no- / explicit on) was verified in an isolated argparse check; full --help needs sglang locally
• The new flag appears in the CLI reference docs
• No example added (N/A)

🤖 Generated with Claude Code