[FSDP][6/N] Check valid param freezing for ModuleWrapPolicy

[awgu]

Stack from ghstack (oldest at bottom):

• [FSDP][9/N] Introduce CustomPolicy #104986
• [FSDP][8/N] Replace _FSDPPolicy.policy with _Policy._run_policy #104969
• [FSDP][Easy] Allow ModuleWrapPolicy to take Iterable #104999
• [FSDP][7/N] Add warning about frozen params #104967
• -> [FSDP][6/N] Check valid param freezing for ModuleWrapPolicy #104427

This PR adds improved error/warning messaging when auto wrapping with ModuleWrapPolicy in the presence of frozen parameters.

• For use_orig_params=False, FSDP requires uniform requires_grad for each FSDP instance. This PR adds a ValueError at wrapping time with a message that mentions the violating module and the frozen/non-frozen parameter names.
• For use_orig_params=True, FSDP allows non-uniform requires_grad for each FSDP instance. However, it will result in higher-than-expected gradient memory usage. This PR adds a UserWarning at wrapping time with a message that mentions the violating module, how much extra gradient memory will be used (in units of numel), and the frozen/non-frozen parameter names.
  • There is a possibility that this warning will be spammy/verbose, but my current thinking is that it is okay for now unless users complain.

Why DFS via named_children() vs. Using named_modules()

LoraModel(
  (embed_tokens): Embedding(100, 32)
  (layers): ModuleList(
    (0-3): 4 x LoraDecoder(
      (attn): LoraAttention(
        (q_proj): Linear(in_features=32, out_features=32, bias=False)
        (lora_A): Linear(in_features=32, out_features=8, bias=False)
        (lora_B): Linear(in_features=8, out_features=32, bias=False)
        (k_proj): Linear(in_features=32, out_features=32, bias=False)
        (v_proj): Linear(in_features=32, out_features=32, bias=False)
        (o_proj): Linear(in_features=32, out_features=32, bias=False)
      )
      (mlp): LoraMLP(
        (proj1): Linear(in_features=32, out_features=128, bias=False)
        (proj2): Linear(in_features=128, out_features=32, bias=False)
      )
      (inp_layernorm): LayerNorm((32,), eps=1e-05, elementwise_affine=True)
      (post_attn_layernorm): LayerNorm((32,), eps=1e-05, elementwise_affine=True)
    )
  )
  (norm): LayerNorm((32,), eps=1e-05, elementwise_affine=True)
)

Reverse topological order with stack-based DFS via named_children():

[
  'embed_tokens',
  'layers.0.attn.q_proj', 'layers.0.attn.lora_A', 'layers.0.attn.lora_B', 'layers.0.attn.k_proj', 'layers.0.attn.v_proj', 'layers.0.attn.o_proj', 'layers.0.attn', 'layers.0.mlp.proj1', 'layers.0.mlp.proj2', 'layers.0.mlp', 'layers.0.inp_layernorm', 'layers.0.post_attn_layernorm', 'layers.0',
  'layers.1.attn.q_proj', 'layers.1.attn.lora_A', 'layers.1.attn.lora_B', 'layers.1.attn.k_proj', 'layers.1.attn.v_proj', 'layers.1.attn.o_proj', 'layers.1.attn', 'layers.1.mlp.proj1', 'layers.1.mlp.proj2', 'layers.1.mlp', 'layers.1.inp_layernorm', 'layers.1.post_attn_layernorm', 'layers.1',
  'layers.2.attn.q_proj', 'layers.2.attn.lora_A', 'layers.2.attn.lora_B', 'layers.2.attn.k_proj', 'layers.2.attn.v_proj', 'layers.2.attn.o_proj', 'layers.2.attn', 'layers.2.mlp.proj1', 'layers.2.mlp.proj2', 'layers.2.mlp', 'layers.2.inp_layernorm', 'layers.2.post_attn_layernorm', 'layers.2',
  'layers.3.attn.q_proj', 'layers.3.attn.lora_A', 'layers.3.attn.lora_B', 'layers.3.attn.k_proj', 'layers.3.attn.v_proj', 'layers.3.attn.o_proj', 'layers.3.attn', 'layers.3.mlp.proj1', 'layers.3.mlp.proj2', 'layers.3.mlp', 'layers.3.inp_layernorm', 'layers.3.post_attn_layernorm', 'layers.3',
  'layers', 'norm', ''
]

Reverse topological order with named_modules():

[
  'norm',
  'layers.3.post_attn_layernorm', 'layers.3.inp_layernorm', 'layers.3.mlp.proj2', 'layers.3.mlp.proj1', 'layers.3.mlp', 'layers.3.attn.o_proj', 'layers.3.attn.v_proj', 'layers.3.attn.k_proj', 'layers.3.attn.lora_B', 'layers.3.attn.lora_A', 'layers.3.attn.q_proj', 'layers.3.attn', 'layers.3',
  'layers.2.post_attn_layernorm', 'layers.2.inp_layernorm', 'layers.2.mlp.proj2', 'layers.2.mlp.proj1', 'layers.2.mlp', 'layers.2.attn.o_proj', 'layers.2.attn.v_proj', 'layers.2.attn.k_proj', 'layers.2.attn.lora_B', 'layers.2.attn.lora_A', 'layers.2.attn.q_proj', 'layers.2.attn', 'layers.2',
  'layers.1.post_attn_layernorm', 'layers.1.inp_layernorm', 'layers.1.mlp.proj2', 'layers.1.mlp.proj1', 'layers.1.mlp', 'layers.1.attn.o_proj', 'layers.1.attn.v_proj', 'layers.1.attn.k_proj', 'layers.1.attn.lora_B', 'layers.1.attn.lora_A', 'layers.1.attn.q_proj', 'layers.1.attn', 'layers.1', 'layers.0.post_attn_layernorm', 'layers.0.inp_layernorm', 'layers.0.mlp.proj2', 'layers.0.mlp.proj1', 'layers.0.mlp', 'layers.0.attn.o_proj', 'layers.0.attn.v_proj', 'layers.0.attn.k_proj', 'layers.0.attn.lora_B', 'layers.0.attn.lora_A', 'layers.0.attn.q_proj', 'layers.0.attn', 'layers.0',
  'layers', 'embed_tokens', ''
]

With the stack-based DFS via named_children(), reversing the topological order gives us each level in the module tree in the registered order, wheres with named_modules(), reversing the topological order gives us each level in reverse. Both are valid orders, but we prefer the former since it allows us to error/warn on the first-registered module that violates the frozen/non-frozen condition.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/104427

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit a03137f:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ezyang]

This function really smacks of "we want to know what the module tree would look like after we do wrapping, but we need to run this before we actually start wrapping, so we have to simulate what the effect would be." That's annoying, because if you change how the wrap works, you also have to adjust this function too (speaking of which, if this statement is true, it would be helpful to have some invariant, like for example _get_param_to_fqn(post_wrapping_module) should report the same as _get_managed_param_to_fqn(pre_wrapping_module)--I'm not sure if this holds, but it would be helpful to know if it is.)

[ezyang]

We are assuming no weight tying, is that right?

[awgu]

We should not be making that assumption. However, for shared parameters, FSDP needs them to be assigned to at least the lowest common ancestor module or higher. A follow-up work item is to add this check to the auto wrapping path.

[ezyang]

instead of reporting numel, reporting nbytes seems better since that will adjust for differing dtype size (...well, unless there's some shenanigans where dtypes have to get bucketed together?)

[awgu]

For fp8, we will need to add support for mixing dtypes in the same flat parameter. It should still be possible to pre-compute the bytes for these tensors though. It just requires some more plumbing.

[ezyang]

In general, this requires splitting your module into multiple modules, some of which have frozen params and some of which don't, right? Is it always safe to take a module M and create two submodules beneath it M1 and M2, one of which has frozen params and the other which doesn't (even if I don't call M1.forward/M2.forward?)

[awgu]

In general, this requires splitting your module into multiple modules, some of which have frozen params and some of which don't, right?

Yes.

Is it always safe to take a module M and create two submodules beneath it M1 and M2, one of which has frozen params and the other which doesn't (even if I don't call M1.forward/M2.forward?)

Unfortunately, it is important to still call M1.forward/M2.forward in order to run pre/post-forward logic, which would notably include the corresponding all-gathers in the pre-forward. In other words, if we just organize the parameters into modules but do not change the parent's forward to actually call those modules' forward, then they will not have their parameters all-gathered.

[ezyang]

It's non-obvious to me why not use_orig_params doesn't work. Without using orig params, we are reconstructing the params from the flat buffer every iteration. What is the requires_grad on this flat buffer? Clearly it has to be requires_grad=True, since we want to accumulate gradients into it. So how do I reconstruct a parameter that doesn't require grad from this requires grad parameter? The naive thing to do is detach (letting me get at the data without setting up a gradient edge.) Is there something in the reduce-scatter that chokes afterwards?

[awgu]

For use_orig_params=True, mixing frozen and non-frozen does work functionally. It just may use memory higher than expected (hence a warning instead of an error).

For the frozen original parameters, we can just have them view into the flat parameter but still have requires_grad=False. FSDP takes care to not forward a gradient into the .grad attribute for those parameters.

[ezyang]

Yes, the use_orig_params=True case makes sense to me. I'm wondering about use_orig_params=False haha

[awgu]

Ahh, I somehow missed the not in not use_orig_params.

I think the issue is for the optimizer. For use_orig_params=False, the optimizer sees the sharded FlatParameters themselves. If the FlatParameter logically contains both frozen and non-frozen parameters, one thing we can do is enforce zeros in all gradient elements for frozen parameters part of the FlatParameter; however, this assumes that for the optimizer, a step with zero gradient is the same as no step (which I think is not always true). Nonetheless, maybe it is not unreasonable to just follow this behavior and set the flat_param.requires_grad = any(orig_param.requires_grad for orig_param in flat_param._params).

[ezyang]

Ah, the interaction with optimizer certainly makes sense.

That being said, most optimizers, if you only ever have steps with zero gradients, it will be equivalent to not updating (e.g., momentum will always be zero). But you are losing most of the benefit from freezing in the first place here, so this is mostly an academic discussion.

[ezyang]

nit: consider using assertExpectedInline! It makes updating these tests if you tweak the warning/error message much more pleasant, since you can just run the test with EXPECTTEST_ACCEPT=1 and it will automatically update all the assertions with the new text.

[ezyang]

Nice! Everything algorithmically looks good. The main annoyance is logic duplication when you have to simulate the post-order wrapping process to accurately know which parameters bucket together. Will have to simultaneously update this code as well as the FSDP code if we ever change it... not that there really is much way we can change it, I guess 🤔

[awgu]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here