[dyurk-lila]

What

The non-packing path of SFTTrainer._validate_batch_parallelism computed the Megatron DP size as total_gpus // (tp * pp), which leaves out context parallelism. It now calls the existing self._dp_size(), which divides by tp * pp * cp. The packing path already uses self._dp_size(). The per_dp_batch / micro-batch check right after used the same wrong value, so it is fixed too.

Why

Megatron's DP size is world / (tp * pp * cp). With CP > 1, the old check required the batch size to be divisible by a number that was too large. For example, 16 GPUs at TP2/CP8 has a real DP size of 1, but batch_size=2 was rejected:

ValueError: batch_size (2) must be divisible by data-parallel size (8)

Tests

Added test_validate_batch_parallelism_counts_context_parallelism to tests/train/sft/test_sft_trainer.py. It is a CPU-only test that builds a bare trainer holding only the config the check reads. It covers:

• 16 GPUs at TP2/CP8 (DP=1): batch_size=2 is accepted. On the old code it fails with the error above.
• 16 GPUs at TP2/CP2 (DP=4): batch_size=4 is accepted, and batch_size=2 is still rejected with "data-parallel size (4)".

pytest tests/train/sft/test_sft_trainer.py: 3 passed. ruff (v0.11.9) and black (24.10.0) are clean.

🤖 Generated with Claude Code

---

Note

Low Risk
Narrows a startup validation bug for Megatron CP; behavior change only affects configs that were incorrectly rejected before.

Overview
Fixes Megatron SFT startup validation so data-parallel size includes context parallelism, matching Megatron’s real DP group size.

In _validate_batch_parallelism, the non–sequence-packing Megatron path no longer uses total_gpus // (tp * pp); it calls self._dp_size() (tp * pp * cp). That same DP value drives the per_dp_batch / micro-batch divisibility check, so configs with CP > 1 are not rejected with an inflated DP size (e.g. batch_size=2 on 16 GPUs at TP2/CP8).

Adds a CPU-only test that exercises accept/reject cases for TP/CP combinations.

^{Reviewed by Cursor Bugbot for commit 5953ecf. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request refactors the _validate_batch_parallelism method in SFTTrainer to use the internal _dp_size() helper method instead of manually calculating the data-parallel size. This ensures context parallelism is properly accounted for. Unit tests have also been added to verify the validation logic under various parallel configurations. There are no review comments to address, and we have no further feedback to provide.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] This PR appears safe to merge.

_{Reviews (1) · Last reviewed commit: "[sft] Count context parallelism in the b..." · Reviewed by Greptile}