feat: add context parallel support for SFT

[ashors1]

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Changelog

• Please update the CHANGELOG.md under next version with high level changes in this PR.

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation? Make sure to also update the NeMo Framework User Guide which contains the tutorials

Checklist when contributing a new algorithm

• Does the trainer resume and restore model state all states?
• Does the trainer support all parallelism techniques(PP, TP, DP)?
• Does the trainer support max_steps=-1 and validation?
• Does the trainer only call APIs defined in alignable_interface.py?
• Does the trainer have proper logging?

Additional Information

• Related to # (issue)

[ashors1]

Note that CP support for SFT was recently added to NeMo. We need a NeMo commit at least as recent as 8c921dc19a905d8b5a0f90f6e2a34607c2e0660d

[terrykong]

TODO:

• potentially rebase after https://github.com/NVIDIA/NeMo-Aligner/actions/runs/12149989736/job/33882060795?pr=405#step:5:743
  • presubmit didn't pass for SFT on that PR, so if that doesn't get resolved by thursday, let's just merge this as is assuming CI passes since #405 should be able to rebase on this easily

[terrykong]

For my education, why do we need to specify this now?

[ashors1]

I don't actually think it's necessary since TE is enabled by default. I just wanted to make explicit the fact that we were using TE. But I will remove this

[terrykong]

According to that fp8 comment above, should this be:

Suggested change

	pad_seq_length_to_mult = 16
	if model_cfg is not None:
	pad_seq_length_to_mult = (
	8 * model_cfg.get("tensor_model_parallel_size", 1) if model_cfg.get("sequence_parallel", False) else 16
	)
	pad_seq_length_to_mult *= model_cfg.get("context_parallel_size", 1)
	pad_seq_length_to_mult = 16
	if model_cfg is not None:
	if model_cfg.get("sequence_parallel", False):
	pad_seq_length_to_mult = math.lcm(pad_seq_length_to_mult, model_cfg.get("tensor_model_parallel_size", 1))
	pad_seq_length_to_mult *= model_cfg.get("context_parallel_size", 1)

? From the comment it sounds like if someone is doing fp8 SFT with TP=1 and set sequence_parallel, then the padding would be too small

[ashors1]

That chunk was taken directly from here: https://github.com/NVIDIA/NeMo/blob/b847bf75c371931e4f17ea426741c1d023afa0c0/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py#L262-L268, but the code does seem to contradict the comment. I'll follow up with the TE team

[ashors1]

From the TE team: "when SP=True, the dimensions are flipped so the sequence dimension is first. So we only need to make sure it's divisible by 8 after the TP split to comply with TE's expectations."

[terrykong]

what does tuning do?

[ashors1]

It controls the keys that are returned in the batch: https://github.com/NVIDIA/NeMo/blob/b847bf75c371931e4f17ea426741c1d023afa0c0/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py#L1211-L1222. If tuning=False, we don't return the keys that are necessary for sequence packing (and thus CP) in TE.

Note also that tuning is set to True in NeMo's SFT: https://github.com/NVIDIA/NeMo/blob/b847bf75c371931e4f17ea426741c1d023afa0c0/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py#L407

[ashors1]

closing in favor of #430. @terrykong I'll address your comments there