Fix FSDP hanging for multimodal models

[jinsooihm]

Problem

Consider a mutlimodal model that only runs its vision encoder when an image is present in the input.
Suppose DP = 2 and rank 0 gets an image rank 1 does not.
With FSDP on the vision encoder, rank 0 will issue an all gather, which will be not met because rank 1 does not run its vision encoder.

This problem exists in the current multimodal models, but the cc12m dataset does not expose it because its samples are always text + image.

I have only implemented this for Muse Glimmer for now but it can be extended to other models if you agree with the design.

Proposed solution

1. Run vision encoder with always and if input is None, return an empty tensor, so all gather is met.
2. Add a zero value related to the empty tensor on the embedded tokens to make encoder backward execute, since there is a reduce scatter that needs to be met.
3. This requires FSDP set_reduce_scatter_unused_params(True) which sets grad to 0 for unused parameters.
   The above 3 steps solve the hanging.

Alternate design

An alternate design could be to wrap the whole vision region + scattering as a separate module vision_tower and let vision_tower.forward(embedded_tokens, pixel_values, ...) that

• if pixel_values is None returns embedded_tokens right away
• run the encoder and scatter and returns embedded_tokens_with_vision_features.

We can then wrap vision_tower as a whole fsdp region. This can remove the hacky stuff at step 2 in the proposed solution.

I think this could be cleaner, but can change FQN of the vision parts.

There is another subtlety coming from optimization. If there are many optimization steps without an image, and we use AdamW, the weight of the vision encoder will keep decaying. To prevent that, we

1. Record whether the model got an image in the input in preprocess_inputs
2. Register a pre optimizer hook that sets gradient of unused parameters to be None if no DP rank saw an image in its optimization step, including PP and grad accumulation microbatches. This requires communication and D2H sync between backward and optimizer step.
3. Then the optimizer does not do optimizer step to those parameters.

This may be a data problem and decay is the intended behavior? Then the solution can be simple by just doing the first 3 steps.
Please comment, cc. @shuhuayu @tianyu-l

Limitations

This makes optimizer step not cudagraphable.
One way to get around this would be to have to separate optimizers for the decoder and encoder part, and run one conditionally, but that would be more invasive change.

[tianyu-l]

why don't we always turn this on -- is there a case in torchtitan where it does unecessary reduction?

asking from another angle: what if vision encoder is frozen, do we set this to True?

[fegin]

set_reduce_scatter_unused_params only performs reduce_scatter when requires_grad is True. So, it seems that we should always turn on this flag.

[tianyu-l]

@weifengpy do you know why it's not True by default in FSDP?

[weifengpy]

this regresses memories. set_reduce_scatter_unused_params is a lazy workaround for non-spmd reduce-scatter. for this case, it's related to some rank didn't trigger RS from vision modules, while other ranks trigger

non-lazy way is to address the root cause of non-spmd vision behavior direclty

[weifengpy]

numerics are also problematic - some ranks just do zero-valued gradients. but those gradients are averaged with real gradients from other ranks

so address non-spmd vision behavior is more rigor than using this API

[shuhuayu]

suppose some ranks get vision input, some ranks do not, then i think we should average gradients by treating the current dp rank's vit's gradient as zero?

[weifengpy]

suppose some ranks get vision input, some ranks do not, then i think we should average gradients by treating the current dp rank's vit's gradient as zero?

after grad reduction, all real gradients shrinks towards zero a little bit

[shuhuayu]

after grad reduction, all real gradients shrinks towards zero a little bit

my understanding is that if we consider the vit + decoder as one single model, and lm loss is defined over all valid tokens, then we should be consistent while scaling the loss/gradient, i.e. using number of valid tokens, whether there are images contributing vit gradients locally is not relevant.

[tianyu-l]

numerics are also problematic - some ranks just do zero-valued gradients. but those gradients are averaged with real gradients from other ranks

We discussed this offline. I think this is expected. This sounds more of a data distribution characteristic / problem, rather than infra problem.

I agree that we should keep spmd behavior as much as possible.

[weifengpy]

got you. spmd is the main reason then