Remove activation checkpointing tag to get correct FQNs (#124698)

[mvpatel2000]

Cherry-pick for release branch

Fixes #124546

When setting use_orig_params = False and using activation checkpointing, the FQN mapping as retrieved by the _get_fqns function is incorrect because the prefix that is added to the name of each activation checkpointed module, _checkpoint_wrapped_module, can still be present. I think this is an edge case with the _get_fqns function that was not addressed by this previous commit #118119.

Without the change, the list of object names for an activation checkpointed module with FSDP (and use_orig_params=False) can be something like:

['model', '_fsdp_wrapped_module', 'transformer', 'blocks', '0', '_fsdp_wrapped_module', '_checkpoint_wrapped_module', '_flat_param']

Which will incorrectly return just one FQN, {'model.transformer.blocks.0._flat_param'}, when all the FQNs of the parameters of the transformer block should be returned.

With the change, the list of object names will now have _checkpoint_wrapped_module removed:

['model', '_fsdp_wrapped_module', 'transformer', 'blocks', '0', '_fsdp_wrapped_module', '_flat_param']

And the FQNs are correctly retrieved and returned in _get_fqns when this condition is satisfied. The correct FQNs are:

{'model.transformer.blocks.0.attn.Wqkv.bias', 'model.transformer.blocks.0.ffn.up_proj.bias',
'model.transformer.blocks.0.attn.out_proj.weight', 'model.transformer.blocks.0.norm_2.weight',
'model.transformer.blocks.0.ffn.down_proj.weight', 'model.transformer.blocks.0.attn.Wqkv.weight',
'model.transformer.blocks.0.norm_2.bias', 'model.transformer.blocks.0.ffn.up_proj.weight',
'model.transformer.blocks.0.ffn.down_proj.bias', 'model.transformer.blocks.0.norm_1.bias',
'model.transformer.blocks.0.norm_1.weight', 'model.transformer.blocks.0.attn.out_proj.bias'}

Pull Request resolved: #124698
Approved by: https://github.com/Skylion007

Fixes #ISSUE_NUMBER

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k @LucasLLC

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit e4fe159 with merge base 86a2d67 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / linux-focal-py3.12-clang10 / test (default, 1, 3, linux.2xlarge) (gh) (disabled by #123655 but the issue was closed recently and a rebase is needed to make it pass)
  test_binary_ufuncs.py::TestBinaryUfuncsCPU::test_logaddexp_cpu_complex128

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