Post

Megatron-LM Fine-Tuning Deep Dive: Qwen Recipes, Datasets, and the Training Loop

A code-level tour of Megatron fine-tuning: where Qwen recipes live, how pretraining and SFT datasets are built, how pretrain_gpt.py reaches the forward/backward schedule, which extension points matter, and a successful eight-GPU Qwen-style MCore smoke test with TP, PP, CP, and EP all set to one.

Posted
By
10 min read
Megatron-LM Fine-Tuning Deep Dive: Qwen Recipes, Datasets, and the Training Loop

Why this post

Megatron is often introduced as a collection of parallelism techniques:

  • tensor parallelism;
  • pipeline parallelism;
  • context parallelism;
  • expert parallelism;
  • distributed optimizer and data parallelism.

That description is correct, but it is not the best place to start if the task is fine-tuning a model.

For fine-tuning, the more useful questions are:

  1. Where is the model architecture selected?
  2. What does the dataset have to look like?
  3. Which function owns the training loop?
  4. Where can I replace the model, forward step, loss, optimizer, or data source?
  5. Can Megatron run without TP, PP, CP, or EP?

This post answers those questions by following the open-source code.

All Megatron-LM links are pinned to commit 78901d8a. Qwen-specific recipe links are pinned separately to Megatron-Bridge commit fcbb6031.

1. Megatron-LM and Megatron-Bridge are different layers

The first important distinction is repository scope.

Megatron-LM

NVIDIA/Megatron-LM contains:

  • Megatron Core model building blocks under megatron/core;
  • the high-level pretrain_gpt.py entrypoint;
  • distributed initialization;
  • dataset builders;
  • pipeline schedules;
  • optimizer and checkpoint infrastructure.

Its GPTModel is architecture-configurable. It is not organized as one Python model class per Hugging Face model family.

Megatron-Bridge

NVIDIA-NeMo/Megatron-Bridge adds the model-specific and Hugging Face-facing layer:

  • Qwen, Llama, DeepSeek, and other recipes;
  • Hugging Face to Megatron conversion;
  • pretraining, SFT, and PEFT configurations;
  • higher-level training configuration containers.

This matters because Megatron-LM itself has no dedicated examples/qwen directory. The official Qwen recipes live in Megatron-Bridge.

For example, the Bridge recipe qwen3_600m_pretrain_config() constructs a Qwen3 0.6B provider and explicitly sets:

1
2
3
4

cfg.model.tensor_model_parallel_size = 1
cfg.model.pipeline_model_parallel_size = 1
cfg.model.context_parallel_size = 1
cfg.model.sequence_parallel = False

The corresponding qwen3_600m_sft_config() is also designed for TP=1 and PP=1 on one node.

So Megatron does not require model parallelism. If every model-parallel size is one, eight GPUs simply form an eight-rank data-parallel group.

2. What a Qwen fine-tuning case looks like

A supervised fine-tuning run has three pieces:

1
2
3

pretrained checkpoint
  + chat-style dataset
  + the same causal language-model training loop

The Qwen Bridge SFT recipe configures a Qwen model provider, tokenizer, sequence length, batch size, optimizer precision, DDP behavior, and checkpoint locations. The recipe leaves a clear checkpoint extension point:

1

# cfg.checkpoint.pretrained_checkpoint = "/path/to/checkpoint"

The distinction between pretraining and fine-tuning is therefore not a separate distributed engine. It is primarily:

  • which checkpoint initializes the model;
  • which dataset class produces tokens and labels;
  • which tokens are included in the loss mask;
  • which optimization schedule is selected.

Underneath the recipe layer, both cases eventually reach the same Megatron forward/backward scheduler.

3. Dataset creation: pretraining versus SFT

Megatron supports two very different data paths.

3.1 Pretraining data: JSONL to .bin and .idx

The classic pretraining pipeline starts with JSONL:

1

{"text": "A document to tokenize and train on."}

The official preprocessing entrypoint is tools/preprocess_data.py. Its Encoder parses one JSON object, tokenizes selected keys, optionally appends EOD, and writes Megatron’s indexed dataset format.

A typical command is:

1
2
3
4
5
6
7

python tools/preprocess_data.py \
  --input corpus.jsonl \
  --output-prefix /data/qwen_corpus \
  --tokenizer-type HuggingFaceTokenizer \
  --tokenizer-model Qwen/Qwen3-0.6B \
  --append-eod \
  --workers 8

The result is normally:

1
2

/data/qwen_corpus_text_document.bin
/data/qwen_corpus_text_document.idx

Training then points --data-path at that prefix.

The indexed representation is useful for large pretraining corpora because it supports memory mapping, deterministic sample construction, and dataset blending without repeatedly parsing JSON.

3.2 SFT data: chat JSONL directly

Megatron-LM also has a direct SFT JSONL path. The format is documented by SFTLowLevelDataset:

1
2
3
4
5
6
7

{
  "messages": [
    {"role": "system", "content": "You are a careful math tutor."},
    {"role": "user", "content": "What is 48 + 24?"},
    {"role": "assistant", "content": "72"}
  ]
}

This path does not require .bin/.idx preprocessing. It uses the Hugging Face datasets JSON loader, and SFTDataset.__getitem__ does the important work:

  1. split packed conversations;
  2. call tokenizer.tokenize_conversation(...);
  3. build tokens and assistant targets;
  4. mask non-assistant targets with IGNORE_INDEX;
  5. truncate or pad to the configured sequence length;
  6. produce packed-sequence metadata such as cu_seqlens.

For my pod experiment, I converted 32 rows from the mounted GSM8K parquet into this format:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

import json
import pyarrow.parquet as pq

rows = pq.read_table(
    "/shared/public/data/gsm8k/train.parquet",
    columns=["extra_info"],
).slice(0, 32).to_pylist()

with open("gsm8k_megatron_sft_sample.jsonl", "w") as output:
    for row in rows:
        info = row["extra_info"]
        json.dump(
            {
                "messages": [
                    {"role": "system", "content": "You are a careful math tutor."},
                    {"role": "user", "content": info["question"]},
                    {"role": "assistant", "content": info["answer"]},
                ]
            },
            output,
        )
        output.write("\n")

The resulting file was:

1

/home/jobuser/gsm8k_megatron_sft_sample.jsonl

with 32 conversations.

3.3 How the dataset class is selected

The high-level selection happens in train_valid_test_datasets_provider:

1
2
3
4
5
6
7
8

if args.sft:
    dataset_type = SFTDataset
elif args.mock_data:
    dataset_type = MockGPTDataset
elif args.fim_data:
    dataset_type = GPTFIMDataset
else:
    dataset_type = GPTDataset

That function is itself an extension point: a different callback can be passed to pretrain(...) to supply entirely different train, validation, and test datasets.

4. From pretrain_gpt.py to one optimizer step

The public training entrypoint is pretrain_gpt.py.

At a high level:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

pretrain_gpt.py
  -> parse and validate arguments
  -> build GPT/model configuration
  -> call pretrain(...)
  -> initialize distributed state
  -> build model and optimizer
  -> build datasets and iterators
  -> enter train(...)
  -> train_step(...)
  -> select forward/backward schedule
  -> forward_step(...)
  -> model(...)
  -> loss_func(...)
  -> finalize gradients
  -> optimizer.step()

The reusable high-level loop starts at training.pretrain. The steady-state loop is train, and one optimizer iteration is orchestrated by train_step.

4.1 The forward-step contract

Megatron does not hard-code the user model call into the distributed scheduler. Instead, the scheduler receives a callback.

The GPT callback is forward_step. It:

  1. obtains a batch;
  2. derives packed-sequence parameters when needed;
  3. calls the model;
  4. returns the model output plus a partially bound loss function.

Conceptually:

1
2
3
4
5
6
7
8
9

output = model(
    tokens,
    position_ids,
    attention_mask,
    labels=labels,
    packed_seq_params=packed_seq_params,
)

return output, partial(loss_func, loss_mask, model=model)

The loss implementation is loss_func. The ordinary language-model branch masks token losses and returns:

  • the local loss sum;
  • the number of valid tokens;
  • reporting metrics.

4.2 Schedule selection

train_step calls get_forward_backward_func(). That factory chooses the schedule based on pipeline topology:

  • no pipeline parallelism: ordinary forward/backward;
  • pipeline parallelism: warmup, steady-state 1F1B, cooldown;
  • virtual pipeline stages: interleaved schedule.

The callback contract stays the same. Changing PP changes the schedule, not the model-facing forward_step API.

5. The main extension points

Megatron is easier to understand when viewed as dependency injection around a distributed loop.

5.1 Model construction

The top-level factory is model_provider. It delegates to a model builder.

The default GPT builder is gpt_builder, which selects the transformer layer specification and creates GPTModel.

Useful model extension points include:

  • a custom model builder;
  • a custom TransformerLayerSpec;
  • Transformer Engine versus local layer implementations;
  • hybrid-model builders;
  • custom embedding and output layers.

5.2 Dataset provider

Replace train_valid_test_datasets_provider to inject a different dataset family, sampling strategy, or storage backend.

The callback returns:

1

(train_dataset, validation_dataset, test_dataset)

5.3 Forward and loss

Replace the forward_step callback to change:

  • batch transformation;
  • multimodal inputs;
  • auxiliary outputs;
  • model call signature;
  • loss closure.

Replace the loss closure to implement distillation, multiple objectives, custom token weighting, or task-specific reporting.

5.4 Optimizer

The optimizer factory is get_megatron_optimizer.

This layer selects among ordinary and distributed optimizer implementations, precision behavior, parameter groups, and optimizer wrappers.

5.5 Parallel topology

The core distributed extension point is topology:

1

world size = DP x TP x PP x CP x EP

Important command-line controls are:

1
2
3
4
5
6
7

--tensor-model-parallel-size
--pipeline-model-parallel-size
--context-parallel-size
--expert-model-parallel-size
--sequence-parallel
--use-distributed-optimizer
--use-megatron-fsdp

Setting TP=PP=CP=EP=1 does not disable Megatron. It leaves model weights unsharded and uses the remaining ranks for data parallelism.

6. Running the official simple loop as a Qwen-style model

Megatron-LM includes examples/run_simple_mcore_train_loop.py. It is the smallest end-to-end example of:

  • distributed initialization;
  • model construction;
  • MockGPTDataset;
  • forward/backward schedule selection;
  • DDP gradient synchronization;
  • optimizer stepping;
  • distributed checkpoint save/load.

The original example uses TP=2 and a toy two-layer model. For the pod smoke test, I kept its loop but changed two things:

  1. model parallelism was set to TP=PP=CP=EP=1;
  2. model_provider() was changed to a Qwen2.5-0.5B-style architecture.

The relevant Qwen-style configuration was:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

config = TransformerConfig(
    num_layers=24,
    hidden_size=896,
    num_attention_heads=14,
    num_query_groups=2,
    kv_channels=64,
    ffn_hidden_size=4864,
    normalization="RMSNorm",
    layernorm_epsilon=1e-6,
    gated_linear_unit=True,
    activation_func=torch.nn.functional.silu,
    add_bias_linear=False,
    add_qkv_bias=True,
    hidden_dropout=0.0,
    attention_dropout=0.0,
    bf16=True,
    params_dtype=torch.bfloat16,
    pipeline_dtype=torch.bfloat16,
    transformer_impl="local",
)

model = GPTModel(
    config=config,
    transformer_layer_spec=get_gpt_layer_local_spec(),
    vocab_size=151936,
    max_sequence_length=128,
    position_embedding_type="rope",
    rotary_base=1000000,
    share_embeddings_and_output_weights=False,
)

This is a random-initialized Qwen-style smoke model, not a converted pretrained Qwen checkpoint. Loading actual Hugging Face Qwen weights is the job of the Bridge/conversion layer.

The distributed setup was:

1
2
3
4
5
6

parallel_state.initialize_model_parallel(
    tensor_model_parallel_size=1,
    pipeline_model_parallel_size=1,
    context_parallel_size=1,
    expert_model_parallel_size=1,
)

and the launch command was:

1
2
3
4

CUDA_DEVICE_MAX_CONNECTIONS=1 \
PYTHONPATH=/home/jobuser/Megatron-LM-78901d8a:$PYTHONPATH \
torchrun --standalone --nproc_per_node=8 \
  /home/jobuser/run_simple_qwen_mcore.py

The pinned source tree was required because the installed wheel did not include the dataset helper Makefile used by compile_helpers().

6.1 Observed result

The model ran three forward/backward/optimizer iterations successfully:

1
2
3
4

iteration=0 losses=[{'lm loss': tensor(12.1715, device='cuda:0')}]
iteration=1 losses=[{'lm loss': tensor(12.1171, device='cuda:0')}]
iteration=2 losses=[{'lm loss': tensor(12.1635, device='cuda:0')}]
Qwen-style MCore smoke test completed.

The entire launch, including C++ dataset-helper compilation and process startup, completed in roughly 32 seconds on eight H100 GPUs.

The important result was not benchmark performance. It was proving that the ordinary Megatron Core loop works with:

1
2
3
4
5

TP=1
PP=1
CP=1
EP=1
DP=8

No model parallelism was necessary.

7. Practical fine-tuning checklist

For a real Qwen SFT job, I would verify the following in order:

  1. Checkpoint conversion
    • Use Megatron-Bridge or another verified converter.
    • Confirm QKV ordering, GQA groups, vocabulary padding, and tied/untied embeddings.
  2. Tokenizer
    • Use the tokenizer corresponding to the checkpoint.
    • Confirm pad, EOD, BOS, and EOS IDs.
    • Verify the chat template used to generate labels.
  3. Dataset
    • Store one messages list per JSONL row.
    • Confirm only assistant tokens contribute to loss.
    • Inspect truncation and packing statistics.
  4. Topology
    • Start with TP=PP=CP=EP=1 if the model fits.
    • Scale model parallelism only when memory or throughput requires it.
  5. Batching
    • Distinguish microbatch size, global batch size, and gradient accumulation.
    • Ensure global batch size is divisible by data-parallel world size.
  6. Numerics
    • Validate BF16/FP8 choices.
    • Compare initial losses before and after checkpoint conversion.
    • Check gradient norms and loss-mask token counts.
  7. Checkpointing
    • Test one save/load round trip before a long run.
    • Confirm optimizer and RNG state behavior for resume.

8. Takeaways

The useful mental model is:

1
2
3

Megatron Core = distributed model and training machinery
Megatron-LM training = reusable pretrain/train/train_step orchestration
Megatron-Bridge = model-specific recipes and HF conversion

For fine-tuning:

  • the SFT dataset determines which tokens carry loss;
  • forward_step and loss_func define the task;
  • the model provider defines the architecture;
  • the optimizer factory defines update behavior;
  • TP, PP, CP, and EP define topology, but each can remain one;
  • data parallelism still gives multi-GPU scaling.

Megatron is not synonymous with model parallelism. It is a modular distributed training system that also works as an ordinary data-parallel trainer when the model fits on one GPU.

References

Machine Learning Training Framework, Distributed Training
megatron-lm megatron-core megatron-bridge qwen fine-tuning sft distributed-training tensor-parallelism pipeline-parallelism dataset pytorch
This post is licensed under CC BY 4.0 by the author.