[FSDP2] Eager-Mode Execution Tracker

[awgu]

Work Items

• Meta-device initialization / _apply() methods
  • Support initial meta-device initialization using swap_tensors path
  • Remove manual padding logic after [Feature][DTensor] Manage additional _padded_local_tensor attribute #113045 @wz337
  • Outcome: Once we have DTensor manage the padded storage, then FSDP only needs to maintain a reference to the DTensor, not its _local_tensor. We should be able to do meta-device initialization via nn.Module.to_empty() (using swap_tensors path) followed by normal initialization ops (e.g. torch.nn.init.<...>) that runs through DTensor's op dispatch without any padding logic in _apply().
• clip_grad_norm_() / sharded gradient scaler
  • Support DTensor-based clip_grad_norm_(), e.g. implementing torch.linalg.vector_norm()
  • Support DTensor-based clip_grad_norm_() with foreach=True
  • Support DTensor-based sharded gradient scaler (low priority since only used for fp16) Gradient scaler for DTensor #132816
  • Outcome: We can replace existing FSDP with per-parameter FSDP in torchtrain.
• 2D sharded state dict
  • Implement strided sharding placement to handle (Shard(0), Shard(0)) placements where FSDP shards dim-0 and TP shards dim-1 [1/N][dtensor] introduce StridedShard placement type and _split_tensor() logic #126697
  • Validate 2D sharded state dict integration with distributed checkpointing excluding checkpoint resharding (so just saving/loading same world size)
  • Validate 2D sharded state dict integration with distributed checkpointing including checkpoint resharding @wz337
  • Outcome: We can do 2D training including checkpointing without communication and can reshard checkpoints to different world sizes.
• distribute_tensor()
  • Allow distribute_tensor() to take in a DTensor to simplify/robust-ify construction of sharded parameter and sharded post-forward parameter @wz337
  • Outcome: We remove all custom DTensor construction code from per-parameter FSDP (_init_sharded_param, _init_sharded_post_forward_param_metadata, to_sharded_post_forward).
• Optimizer <> DTensor
  • Enable foreach=True by default for Adam/AdamW @wz337
  • Support fused=True for Adam/AdamW @wz337
  • Support foreach=True for all torch-native optimizers @wz337
  • Outcome: We have competitive optimizer performance on common optimizers without user-code changes (like passing foreach=True explicitly). This includes new CPU fused Adam/AdamW kernels.
  • Support custom (user-defined) ops [DTensor] add support for custom op registration #131108
  • Outcome: We can support custom optimizers (e.g. Apex, low-precision, etc.).
• Reduce-scatter copy-in kernel/fast-path
  • Add new aten op for chunk-cat @BoyuanFeng
  • Outcome: We can achieve competitive performance with existing FSDP even on low-compute-density workloads (e.g. recommendation models).
• HSDP
  • Add hybrid sharding (when passing 2D mesh arg) @weifengpy
  • Validate HSDP sharded state dict integration with distributed checkpointing
  • Validate HSDP + TP sharded state dict integration with distributed checkpointing
  • Outcome: Existing HSDP has adoption both internally and externally. We should support it under per-parameter and migrate users.
• CPU offloading
  • Add CPU offloading to per-parameter FSDP with async H2D parameter copy and async D2H gradient copy (at least when not accumulating gradient)
• FSDP extensions
  • Add pre/post-all-gather extensions support @awgu
  • Add pre/post-all-gather extensions for float8_experimental @awgu (dynamic scaling eager done)
  • Add pre/post all-gather extensions for QLoRA @weifengpy

References

RFC: #114299

[ad8e]

fp16 is used for diffusion models, so the lack of gradient scaler is a blocker there.

However, #116054 seems relevant, which would complicate work.

[awgu]

Are there any example open source repos for training diffusion models in fp16?

We can try to get some bandwidth to look at DTensor + grad scaler.

[Birch-san]

edm2 trains a diffusion model in fp16 (albeit with some quirks like matmuling with weights which have been normalized in fp32 then downcasted back to fp16). it's a reasonably simple codebase:
https://github.com/Nvlabs/edm2
https://arxiv.org/abs/2312.02696
NOTE: EDM2's float16 training doesn't actually use grad-scaling; they improved the training dynamics to the extent that they were able to turn off grad-scaling.

likewise the original OpenAI UNet (ADM) trained in fp16:
https://github.com/openai/guided-diffusion
https://arxiv.org/abs/2105.05233

We train in 16-bit precision using loss-scaling, but maintain 32-bit weights, EMA, and optimizer state

to be honest, those are actually the only examples I can think of. maybe some parameter-efficient finetunes train in fp16.

HDiT can train for a while in mixed-precision float16, but eventually the colours will fry:
https://github.com/crowsonkb/k-diffusion
but the toy configs (Oxford Flowers, CIFAR-10) should survive for a few minutes of training, if you're just trying to do a test and see "does it get closer to looking like an image".
oh, though I should add that HDiT has a history of "FSDP1 doesn't work on every device we tried it on". bf16 mixed-precision worked on RTX A6000 Adas, but not A40s or 3090s. 3090s worked if mixed-precision was disabled. that was over a year ago, so maybe all sorts of things have changed since then. maybe it was just something stupid like "learned 1-element scalars screwed up the alignment of the flattened weights".

[ad8e]

Update: My own fp16 grad scaler is 50 lines. It doesn't handle the fancy mixed-device things that the PyTorch grad scalar does. It's somewhat faster by avoiding unnecessary kernel launches (filed in another bug) and a CPU-GPU sync (pulling the scale from the GPU each step for logging). So if anyone wants to use fp16, it is easy to roll a good scaler from scratch.

[weifengpy]

Update: My own fp16 grad scaler is 50 lines. It doesn't handle the fancy mixed-device things that the PyTorch grad scalar does. It's somewhat faster by avoiding unnecessary kernel launches (filed in another bug) and a CPU-GPU sync (pulling the scale from the GPU each step for logging). So if anyone wants to use fp16, it is easy to roll a good scaler from scratch.

curious: are you operating directly on DTensor._local_tensor or going through DTensor dispatch ?

[ad8e]

Currently I use to_local() in a no-grad context, then torch._foreach_norm, then torch.stack(), then manually square and all-reduce the norms in one go. There was some issue with calling torch._foreach_norm on DTensor directly, but maybe it was my fault, I didn't diagnose carefully.

I use the per-tensor grad norms for other purposes too, so my format of having all the grad norms in a single tensor fits the other stuff I'm doing.

[weifengpy]

Currently I use to_local() in a no-grad context, then torch._foreach_norm, then torch.stack(), then manually square and all-reduce the norms in one go. There was some issue with calling torch._foreach_norm on DTensor directly, but maybe it was my fault, I didn't diagnose carefully.

I use the per-tensor grad norms for other purposes too, so my format of having all the grad norms in a single tensor fits the other stuff I'm doing.

thanks for the details. Great work!

[ad8e]

I tested _foreach_norm(DTensor); the prior error disappeared. It's slower than using local tensors however. https://api.wandb.ai/links/novelaix/va6xg9g8

Interestingly, the train/overhead panel also shows a slowdown - but this measures time doing Python loop things (wandb, async dataloader, non-syncing if()s), and excludes the optimizer and model. So the expectation is that changing the grad norm calculation shouldn't affect this overhead. Apparently not.

def square_norm_list(list_of_tensors):
    if isinstance(list_of_tensors[0], DTensor):
        list_of_tensors = [p.to_local() for p in list_of_tensors]
    local_norms = torch.stack(torch._foreach_norm(list_of_tensors))
    local_norms.pow_(2)
    dist.all_reduce(local_norms, async_op=False)
    return local_norms

...

    if self.distributed:
        # this is substantially faster than using DTensor dispatch for _foreach_norm
        all_grad_norm_squared = square_norm_list(all_grads)
        all_grad_norm = all_grad_norm_squared.sqrt()
    else:
        all_grad_norm = torch.stack(torch._foreach_norm(all_grads))