feat: Compatibility modification of megatron-fsdp #14593
Conversation
Signed-off-by: Selvaraj Anandaraj <[email protected]> Signed-off-by: jianbinc <[email protected]>
Signed-off-by: jianbinc <[email protected]>
Signed-off-by: jianbinc <[email protected]>
shjwudp
force-pushed
the
jianbinc/nvfsdp_update
branch
from
07e4c4e to
0c17d92
Compare
| @@ -84,6 +84,49 @@ def override_recipe_configs( | |||
| recipe = set_exp_logging_configs( | |||
| recipe, "pre_train", "llm", "llama3", args.tensorboard, args.wandb, args.wandb_prj_name, args.wandb_job_name | |||
| ) | |||
| # for saving checkpoints | |||
| ckpt_path = "/lustre/fsw/coreai_devtech_all/jianbinc/playground/nemo_nvfsdp_update/NeMo/checkpoints" | |||
There was a problem hiding this comment.
We should remove this path
There was a problem hiding this comment.
Yes, thanks for the review!
formt code Signed-off-by: jianbinc <[email protected]>
Signed-off-by: shjwudp <[email protected]>
| @@ -365,9 +365,9 @@ def __init__( | |||
|
|
|||
| self._fsdp = None | |||
|
|
|||
| if fsdp is None and self.ddp_config and self.ddp_config.use_custom_fsdp: | |||
There was a problem hiding this comment.
Customers who built training scripts/tutorials using this would break. We have should comms around it.
Can we should release-note documentation for this ?
There was a problem hiding this comment.
Hi, thanks for point it, I've added use_custom_fsdp as a fallback so that users using older M-Core won't experience breaks.
Can I ask where should I add a release note?
…ore 0.14 Signed-off-by: jianbinc <[email protected]>
… jianbinc/nvfsdp_update
Signed-off-by: shjwudp <[email protected]>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| def set_use_megatron_fsdp(recipe): | ||
| try: | ||
| recipe.trainer.strategy.ddp.use_megatron_fsdp = True | ||
| except: |
Check notice
Code scanning / CodeQL
Except block handles 'BaseException' Note
|
|
||
| try: | ||
| recipe.trainer.strategy.ddp.keep_fp8_transpose_cache = False | ||
| except: |
Check notice
Code scanning / CodeQL
Except block handles 'BaseException' Note
| ) | ||
| try: | ||
| recipe.trainer.strategy.ddp.keep_fp8_transpose_cache = bool(keep_fsdp_fp8_transpose_cache) | ||
| except: |
Check notice
Code scanning / CodeQL
Except block handles 'BaseException' Note
Important
The
Update branchbutton must only be pressed in very rare occassions.An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.
What does this PR do ?
Starting with M-Core v0.14, custom-fsdp was changed to megatron-fsdp, which introduced some API changes. This MR addresses these changes. The most significant change is checkpoint format change from M-Core
torch_disttofsdp_dtensor.Collection: [Note which collection this PR will affect]
Changelog
Usage
# Add a code snippet demonstrating how to use thisGitHub Actions CI
The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.
The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".
Before your PR is "Ready for review"
Pre checks:
PR Type:
If you haven't finished some of the above items you can still open "Draft" PR.
Who can review?
Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.
Additional Information