[training migration] add training config dataclass and arg generation utility

[maanug-nv]

What does this PR do ?

1. Create dataclass for settings related to the training loop.
2. Create an extensible utility to generate argparse arguments from dataclasses

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share discuss a design-doc with the team.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Loading

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

1. Attach the Expert Review label when your PR is ready for review.
2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

1. Add Final Review label
2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either [email protected] or [email protected].

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

[ananthsub]

can you document the argparse_meta metadata schema in more detail in megatron/training/argument_utils.py so it can be used as a reference? otherwise other developers would have to look example usages from these configs

[ananthsub]

nit: might as well use the new style of typehints for the new code

Suggested change

	micro_batch_size: Optional[int] = None
	micro_batch_size: int | None = None

[ananthsub]

exit-signal in the existing argparse is a str, and the conversion to the enum happens here:

Megatron-LM/megatron/training/dist_signal_handler.py

Lines 57 to 59 in f426230

	class DistributedSignalHandler:
	def __init__(self, sig: str = 'SIGTERM'):
	self.sig = SIGNAL_MAP.get(sig, signal.SIGTERM)

[ananthsub]

could you also add a test for a field typed with union to test the fallback on the argparse meta?