# Main Entry Point Script for Training Diffusion Language Models in MegaDLMs

> Discover the main entry point script pretrain_difflm.py for training Diffusion Language Models in MegaDLMs. Access the code and start your training now.

- Repository: [Jinjie Ni/megadlms](https://github.com/jinjieni/megadlms)
- Tags: how-to-guide
- Published: 2026-03-04

---

**The primary entry point for pre-training Diffusion Language Models (Diffusion-LM) in MegaDLMs is the [`pretrain_difflm.py`](https://github.com/jinjieni/megadlms/blob/main/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`](https://github.com/jinjieni/megadlms/blob/main/pretrain_difflm.py) Training Driver

Located at the repository root, [`pretrain_difflm.py`](https://github.com/jinjieni/megadlms/blob/main/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):

```python
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`](https://github.com/jinjieni/megadlms/blob/main/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`](https://github.com/jinjieni/megadlms/blob/main/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`](https://github.com/jinjieni/megadlms/blob/main/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`](https://github.com/jinjieni/megadlms/blob/main/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`](https://github.com/jinjieni/megadlms/blob/main/pretrain_difflm.py) directly with the required Megatron and Diffusion-LM arguments:

```bash
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`](https://github.com/jinjieni/megadlms/blob/main/custom_args/difflm.py):

```bash
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`](https://github.com/jinjieni/megadlms/blob/main/pretrain_difflm.py)** – Main training entry point at the repository root
- **[`custom_args/difflm.py`](https://github.com/jinjieni/megadlms/blob/main/custom_args/difflm.py)** – Diffusion-LM specific command-line arguments
- **[`megatron/core/models/difflm/gpt_model.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/core/models/difflm/gpt_model.py)** – Core `GPTModel` implementation with diffusion-specific forward passes
- **[`megatron/core/models/difflm/gpt_layer_specs.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/core/models/difflm/gpt_layer_specs.py)** – Layer specifications (`get_difflm_layer_*_spec`, `get_difflm_decoder_block_spec`)
- **[`megatron/training/pretrain.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/training/pretrain.py)** – Generic Megatron pre-training driver invoked by the entry script
- **[`megatron/core/datasets/gpt_dataset.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/core/datasets/gpt_dataset.py)** – Dataset class supporting packing for difflm
- **[`tools/checkpoint/convert.py`](https://github.com/jinjieni/megadlms/blob/main/tools/checkpoint/convert.py)** – Checkpoint conversion utilities for difflm models
- **[`examples/dlm_generation/dlm_inference.py`](https://github.com/jinjieni/megadlms/blob/main/examples/dlm_generation/dlm_inference.py)** – Example inference script for trained checkpoints

## Summary

- **[`pretrain_difflm.py`](https://github.com/jinjieni/megadlms/blob/main/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`](https://github.com/jinjieni/megadlms/blob/main/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`](https://github.com/jinjieni/megadlms/blob/main/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`](https://github.com/jinjieni/megadlms/blob/main/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`](https://github.com/jinjieni/megadlms/blob/main/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`](https://github.com/jinjieni/megadlms/blob/main/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`](https://github.com/jinjieni/megadlms/blob/main/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.