Distributed Tensor Parallelism with Megatron-LM and TrainJob

[andreyvelich]

What you would like to be added?

Megatron-LM is a popular library for large-scale distributed training and inference. It implements Tensor Parallelism (TP) and Pipeline Parallelism (PP), enabling the training of LLMs that cannot fit on a single GPU even with FSDP.

We can showcase a Kubernetes-native approach to running Megatron-LM jobs at scale using KF Trainer.

Since Megatron-LM uses torchrun as the primary entrypoint for distributed execution, it would be valuable to provide example Notebook that leverage the torch-distributed runtime. Once @jaiakash finalizes support for GPU ARC runners, we can validate TP Notebook through our E2E tests.

Looking ahead, we can also explore whether introducing a dedicated runtime for Megatron-LM would make sense.

Useful resources:

• Megatron bridge for HF <-> Megatron model conversion: https://github.com/NVIDIA-NeMo/Megatron-Bridge
• https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/quickstart.html
• https://docs.nvidia.com/nemo/megatron-bridge/0.2.0/parallelisms.html#enable-tensor-parallelism)
• [Discussion] PyTorch Operator Improvement #1836

cc @akshaychitneni @kubeflow/kubeflow-trainer-team @kubeflow/kubeflow-sdk-team @bigsur0 @kuizhiqing @zw0610 @johnugeorge @Ronkahn21 @ko3n1g

Why is this needed?

Demonstrate how Trainer can be used for large-scale TP and Megatron-LM

Love this feature?

Give it a 👍 We prioritize the features with most 👍

[andreyvelich]

/good-first-issue

[google-oss-prow]

@andreyvelich:
This request has been marked as suitable for new contributors.

Please ensure the request meets the requirements listed here.

If this request no longer meets these requirements, the label can be removed
by commenting with the /remove-good-first-issue command.

Details

In response to this:

/good-first-issue

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

[XploY04]

@andreyvelich I would like to work on this.

[andreyvelich]

Cool, thanks @XploY04!
/assign @XploY04

[gupta-tilak]

Hi @andreyvelich, I am new to the kubeflow and I am looking to contribute to the kubeflow-trainer. Can you pls assign this to me? My proposal to add distributed TP with Megatron-LM this is mentioned below with reference.

Proposal: Distributed Tensor Parallelism with Megatron-LM and TrainJob

Motivation

Megatron-LM is NVIDIA's widely-adopted library for large-scale distributed training. It implements Tensor Parallelism (TP) and Pipeline Parallelism (PP) — key techniques for training LLMs that exceed single-GPU memory capacity even with FSDP. Since Megatron-LM uses torchrun as its entrypoint, it integrates naturally with Kubeflow Trainer's existing torch-distributed ClusterTrainingRuntime.

There is currently no Megatron-LM example in the project, leaving a gap for users interested in model-parallel training on Kubernetes.

Goals

1. Provide an example Notebook that demonstrates Megatron-Core GPT model training with Tensor Parallelism using the torch-distributed runtime.
2. Follow existing notebook conventions (structure, SDK usage, Papermill parameters for E2E testability).
3. Enable future E2E validation once GPU ARC runners are available.
4. Lay groundwork for evaluating whether a dedicated Megatron-LM runtime is warranted.

Non-Goals

• Introducing a new CRD, runtime, or plugin for Megatron-LM (future work).
• Production-grade training configurations (this is a demonstration).
• Supporting Pipeline Parallelism in the initial example (can be added later).

---

Design

Why torch-distributed Runtime Works

Megatron-LM uses torchrun for launching distributed jobs. The torch-distributed ClusterTrainingRuntime already:

• Sets PET_NNODES, PET_NPROC_PER_NODE, PET_NODE_RANK, PET_MASTER_ADDR, PET_MASTER_PORT environment variables.
• Configures headless services for inter-node communication.
• Supports numProcPerNode as auto, cpu, gpu, or an integer — perfect for mapping to Megatron's tensor_model_parallel_size.

Megatron-Core's parallel_state.initialize_model_parallel(tensor_model_parallel_size) operates within the PyTorch distributed process group that torchrun initializes, so no additional Trainer-side changes are needed.

Example Notebook Design

Location: examples/megatron/tensor-parallelism/megatron-gpt-training.ipynb

Training Script: A self-contained function (following the CustomTrainer(func=...) pattern) that:

1. Initializes PyTorch distributed (via torchrun env vars automatically set by the runtime).
2. Initializes Megatron-Core model parallelism with configurable tensor_model_parallel_size.
3. Creates a small GPT model using megatron.core.models.gpt.GPTModel.
4. Trains for a few iterations on MockGPTDataset (no external data dependency).
5. Saves/loads distributed checkpoints to demonstrate Megatron's checkpoint system.

Notebook Structure (following established conventions from existing examples):

Section	Description
Title + description	Megatron-LM TP training overview, links to docs
Install Kubeflow SDK	pip install -U kubeflow
Define training function	Self-contained megatron_train_gpt() with all imports inside
Scale with TrainJob	TrainerClient() initialization
List runtimes	client.list_runtimes()
Parameters cell	Papermill-injectable num_nodes, num_gpu, tp_size
Submit TrainJob	client.train(trainer=CustomTrainer(...), runtime="torch-distributed")
Check steps	client.get_job(name=...).steps
Watch logs	client.get_job_logs(name, follow=True)
Cleanup	client.delete_job(name)

Papermill Parameters for E2E Testing

The notebook will include a parameters cell compatible with the existing make test-e2e-notebook workflow:

#parameters
num_nodes = 2
num_gpu = 1
tp_size = 2

This enables CI to override values:

make test-e2e-notebook \
  NOTEBOOK_INPUT=./examples/megatron/tensor-parallelism/megatron-gpt-training.ipynb \
  NOTEBOOK_OUTPUT=./artifacts/notebooks/megatron_gpt_training.ipynb \
  PAPERMILL_PARAMS="-p num_nodes 2 -p num_gpu 1 -p tp_size 2" \
  PAPERMILL_TIMEOUT=1800

Container Image Consideration

Megatron-Core requires the megatron-core and NVIDIA's PyTorch NGC runtime with CUDA+NCCL. The notebook will use nvcr.io/nvidia/pytorch:24.02-py3 as the base image (which includes all CUDA/NCCL dependencies), and install megatron-core at runtime via packages_to_install.

Alternatively, the training function installs Megatron-Core inline, keeping the example self-contained without requiring a custom Docker image.

---

Implementation Plan

Phase 1: Example Notebook (this PR)

Step	Task	Files Changed
1	Create example directory examples/megatron/tensor-parallelism/	New directory
2	Implement megatron-gpt-training.ipynb with Megatron-Core TP training	New file
3	Validate notebook structure follows existing patterns	—
4	Update examples/README.md with Megatron-LM section	Existing file

Phase 2: E2E Test Integration (follow-up, after GPU ARC runners)

Step	Task
5	Add Megatron notebook to .github/workflows/test-e2e.yaml
6	Validate TP training on multi-GPU nodes

Phase 3: Dedicated Runtime Evaluation (future)

Step	Task
7	Evaluate if a dedicated megatron-distributed ClusterTrainingRuntime is needed
8	If yes: create runtime YAML, update README.md in manifests

---

Alternatives Considered

1. Using DeepSpeed runtime instead of torch-distributed: Megatron-LM's native distributed backend is torchrun, not MPI-based like DeepSpeed. The torch-distributed runtime is the natural fit since it directly maps to torchrun's environment variable conventions (PET_*).

2. Creating a dedicated Megatron-LM runtime immediately: Premature — the existing torch-distributed runtime already provides everything Megatron-LM needs. A dedicated runtime could be justified later if Megatron-specific defaults (e.g., fixed TP/PP sizes, pre-built NVIDIA images) prove valuable.

3. Using pretrain_gpt.py from Megatron-LM repo: This requires complex CLI argument passing and external data. Using Megatron-Core's Python API with MockGPTDataset keeps the example self-contained and easy to maintain.

---

Testing

• Local: Notebook can be run with LocalProcessBackendConfig on a machine with 2+ GPUs.
• CI/E2E: Papermill parameters allow the notebook to be tested via make test-e2e-notebook once GPU ARC runners are available.
• Manual: The training function uses MockGPTDataset so no external data downloads are needed.

References

• Megatron-Core Quick Start
• Megatron Bridge Parallelisms Guide
• Megatron-LM Repository
• Existing PyTorch DDP Example
• Existing DeepSpeed Example

[XploY04]

Hey @gupta-tilak
I am working on this issue, please help in the review when I raise a PR for it.
Thanks