refactor(fsdp)(4/5): per-parameter weight-precision control for FSDP2 - #91

Closed

Rockdu wants to merge 45 commits into

radixark:mainfrom

Rockdu:feat/fsdp-precision-control

Commits on Aug 3, 2026

feat(fsdp): compile per-parameter precision plans onto FSDP2 with autocast forward
Rockdu
committed
9d5b66a
refactor(fsdp): lower gather overrides to grouped nested fully_shard instead of ignored_params
Rockdu
committed
5af4bfe
perf(fsdp): keep precision sub-shard groups unsharded through backward
Rockdu
committed
987839f
refactor(fsdp): move precision dtype defaults to set_default_diffusion_args and slim compiler output types
Rockdu
committed
d6eb22f
docs(fsdp): comment each stage of compile_precision_plan
Rockdu
committed
926b9a3
docs(fsdp): replace stage comments with a flow diagram in the precision module docstring
Rockdu
committed
ba941fd
style: apply pre-commit isort to precision plan test imports
Rockdu
committed
bb6220b
feat(fsdp): declare model-boundary input dtypes as family policy
Rockdu
committed
0ca24cf
test(e2e): update SD3.5 OCR standard for the autocast forward numerics
Rockdu
committed
1214240

Commits on Aug 4, 2026