Main Entry Point Script for Training Diffusion Language Models in MegaDLMs

The primary entry point for pre-training Diffusion Language Models (Diffusion-LM) in MegaDLMs is the pretrain_difflm.py script located at the repository root.

The MegaDLMs repository extends Megatron-LM to support diffusion-based language modeling architectures. To launch a distributed training run for Diffusion-LM (also referred to as Difflm), practitioners invoke this dedicated driver script that wires together model construction, data loading, and Megatron's distributed training loop.

The pretrain_difflm.py Training Driver

Located at the repository root, pretrain_difflm.py serves as the launching pad for all Diffusion-LM pre-training jobs. The script follows the standard Megatron entry-point pattern, executing the pretrain() function when invoked directly.

When run as a standalone module, the script initializes training through the following guard clause (lines 19–31):

if __name__ == "__main__":
    # ...

    pretrain(
        train_valid_test_datasets_provider,
        model_provider,
        ModelType.encoder_or_decoder,
        forward_step,
        args_defaults={'tokenizer_type': 'HuggingFaceTokenizer'},
        extra_args_provider=extra_args_provider,
    )

This call delegates to megatron.training.pretrain, which handles distributed process spawning, checkpoint management, and the main training loop.

Key Components Orchestrated by the Entry Point

The pretrain_difflm.py script coordinates four critical subsystems through specific provider functions and argument parsers.

Model Construction via model_provider

The model_provider function, defined within pretrain_difflm.py, constructs the Diffusion-LM architecture. It instantiates a GPTModel using diffusion-specific layer specifications imported from the core models directory.

Key implementation details:

• Uses get_difflm_layer_*_spec functions to define transformer block specifications
• Imports GPTModel from megatron/core/models/difflm/gpt_model.py
• Configures the model based on the --model-running-mode argument (e.g., difflm-noshift)

Dataset Provision Pipeline

The script supplies training data through train_valid_test_datasets_provider, which leverages BlendedMegatronDatasetBuilder to create compatible datasets. This provider ensures proper handling of packed token sequences and variable-length masking required for diffusion training.

Diffusion-Specific Arguments

Command-line arguments unique to Diffusion-LM are injected via extra_args_provider, which points to custom_args/difflm.py. This module defines critical hyperparameters including:

• --model-running-mode – Selects the training regime (e.g., difflm-noshift, difflm-noshift-uniform)
• --difflm-varilen-prob – Controls probability of variable-length token masking
• --difflm-time-embedding – Enables optional time-embedding MLP layers

Megatron Infrastructure Integration

The script relies on Megatron's training utilities for distributed orchestration. Functions like get_args, print_rank_0, and the global argument registry manage GPU topology, gradient synchronization, and logging across nodes.

Running a Pre-Training Job

To initiate training, execute pretrain_difflm.py directly with the required Megatron and Diffusion-LM arguments:

python pretrain_difflm.py \
    --model-type difflm \
    --model-running-mode difflm-noshift \
    --micro-batch-size 4 \
    --global-batch-size 256 \
    --seq-length 2048 \
    --max-position-embeddings 2048 \
    --train-iters 100000 \
    --lr 0.0006 \
    --data-path /path/to/megatrons/data

• The --model-type difflm flag selects the Diffusion-LM architecture
• --model-running-mode difflm-noshift selects the standard no-shift training regime

Customizing the Diffusion Schedule

You can modify the diffusion behavior through arguments defined in custom_args/difflm.py:

python pretrain_difflm.py \
    --model-running-mode difflm-noshift-uniform \
    --difflm-varilen-prob 0.02 \
    --difflm-time-embedding \
    --data-path /path/to/data

Core Repository File Architecture

The following files constitute the end-to-end training pipeline for Diffusion Language Models in MegaDLMs:

• pretrain_difflm.py – Main training entry point at the repository root
• custom_args/difflm.py – Diffusion-LM specific command-line arguments
• megatron/core/models/difflm/gpt_model.py – Core GPTModel implementation with diffusion-specific forward passes
• megatron/core/models/difflm/gpt_layer_specs.py – Layer specifications (get_difflm_layer_*_spec, get_difflm_decoder_block_spec)
• megatron/training/pretrain.py – Generic Megatron pre-training driver invoked by the entry script
• megatron/core/datasets/gpt_dataset.py – Dataset class supporting packing for difflm
• tools/checkpoint/convert.py – Checkpoint conversion utilities for difflm models
• examples/dlm_generation/dlm_inference.py – Example inference script for trained checkpoints

Summary

• pretrain_difflm.py is the main entry point script for training Diffusion Language Models in MegaDLMs, located at the repository root.
• The script executes the pretrain() function from Megatron's training library when invoked directly.
• model_provider constructs the GPT-based Diffusion-LM using specifications from megatron/core/models/difflm/.
• extra_args_provider imports diffusion-specific hyperparameters from custom_args/difflm.py.
• The script supports distributed training via Megatron's infrastructure and accepts standard Megatron arguments alongside Diffusion-LM specific flags.

Frequently Asked Questions

Where is the main entry point script for training Diffusion Language Models in MegaDLMs?

The main entry point is pretrain_difflm.py located at the root of the MegaDLMs repository. This script contains the if __name__ == "__main__" guard that launches the Megatron training loop when executed directly with Python.

How do I specify Diffusion-LM specific training arguments?

Diffusion-LM specific arguments are defined in custom_args/difflm.py and injected through the extra_args_provider parameter. You can pass flags like --model-running-mode and --difflm-varilen-prob on the command line when calling pretrain_difflm.py.

What model architecture does the entry point script construct?

The script constructs a GPTModel using diffusion-specific layer specifications via the model_provider function. It imports get_difflm_layer_*_spec functions from megatron/core/models/difflm/gpt_layer_specs.py to build transformer blocks optimized for diffusion language modeling.

Can I use standard Megatron-LM arguments with this script?

Yes, the pretrain_difflm.py script accepts all standard Megatron-LM arguments (such as --micro-batch-size, --global-batch-size, and --data-path) while adding Diffusion-LM specific parameters. It delegates to megatron.training.pretrain for handling distributed training infrastructure.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}