Fix FSDP2 strategy related bugs

[BoxiangW]

No description provided.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Copilot]

Pull Request Overview

This PR fixes known FSDP2 strategy issues by updating distributed configuration settings, refining checkpoint scheduling, and updating dataset loading parameters. Key changes include:

• Adding new distributed configuration options (dp_size, tp_size, cp_size, sequence_parallel) and dataset parameter (num_samples_limit) in the YAML recipe.
• Updating checkpoint scheduling logic to skip checkpointing on step 0.
• Refactoring FSDP2 manager imports and type annotations, and adding a trust_remote_code parameter to the HellaSwag dataset.

Reviewed Changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
recipes/llama_3_2_1b_hellaswag.yaml	Updated distributed and dataset configurations to support FSDP2 fixes.
automodel/training/step_scheduler.py	Modified checkpoint condition to avoid triggering on step 0.
automodel/training/base_recipe.py	Added directory creation for checkpoint saving.
automodel/distributed/parallelizer.py	Removed safe_import_from wrappers in favor of direct imports.
automodel/distributed/fsdp2.py	Updated type annotations and default tensor parallel sharding plan.
automodel/datasets/utils.py	Minor punctuation fixes in comments.
automodel/datasets/llm/hellaswag.py	Updated dataset initialization to include trust_remote_code.

[akoumpa]

Thanks a lot @BoxiangW ,

I fixed one minor issue with the linter and DCO'd.

One thing we may want to do in the future is have a function provide the default TP-plan/s, but I think for now this is mergable and we should proceed -- a future note for us to add in upcoming PRs.

Thank you.