pretrain_diffusion.py

# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

"""Pretrain Masked Diffusion Language Model (MDLM)."""

# Capture the true program start time BEFORE any heavy imports.
import time
_PROGRAM_START_TIME = time.time()

# Suppress warnings on all ranks but rank 0.
import os
import warnings
rank = int(os.environ.get('RANK', 0))
if rank != 0:
    warnings.filterwarnings("ignore", category=UserWarning)
    warnings.filterwarnings("ignore", category=FutureWarning)

from functools import partial
from typing import Optional

import torch

from diffusion_builders import diffusion_builder
from megatron.core import parallel_state
from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset
from megatron.training.datasets.sft_dataset import SFTDataset
from megatron.core.enums import ModelType
from megatron.core.models.diffusion.alpha_scheduler import get_alpha_scheduler
from megatron.core.rerun_state_machine import get_rerun_state_machine
from megatron.core.utils import get_attr_wrapped_model, StragglerDetector
from megatron.training import (
    get_args,
    get_timers,
    get_tokenizer,
    pretrain,
    print_rank_0,
    set_startup_timestamps,
)
from megatron.training.arguments import core_transformer_config_from_args
from megatron.training.utils import (
    get_batch_on_this_cp_rank,
    get_batch_on_this_tp_rank,
    get_blend_and_blend_per_split,
    is_first_or_last_pipeline_stage,
)
from model_provider import model_provider

stimer = StragglerDetector()

# Module-level scheduler cache (instantiated once per process).
_scheduler = None


def _get_scheduler():
    """Return the cached alpha scheduler, instantiating it on first call."""
    global _scheduler
    if _scheduler is None:
        args = get_args()
        scheduler_name = getattr(args, 'diffusion_alpha_scheduler', 'LinearAlphaScheduler')
        _scheduler = get_alpha_scheduler(scheduler_name)
    return _scheduler


def get_batch(data_iterator, vp_stage: Optional[int] = None):
    """Generate a batch.

    Returns tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params.
    Mirrors the structure of pretrain_gpt.py's get_batch().
    """
    args = get_args()
    config = core_transformer_config_from_args(args)

    if not is_first_or_last_pipeline_stage(vp_stage):
        return None, None, None, None, None, None

    # Get batches based on the TP rank.
    batch = get_batch_on_this_tp_rank(data_iterator)

    batch.pop('cu_seqlens', None)
    batch.pop('cu_seqlens_padded', None)
    batch.pop('max_seqlen', None)
    local_cp_size = batch.pop('local_cp_size', None)

    # Slice batch along sequence dimension for context parallelism.
    batch = get_batch_on_this_cp_rank(batch)
    packed_seq_params = None

    return (*batch.values(), packed_seq_params)


def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor, model=None):
    """Diffusion loss function.

    Args:
        loss_mask (torch.Tensor): Diffusion-weighted loss mask [b, s].
            Encodes masking pattern, scheduler weights, and normalization factor.
        output_tensor (torch.Tensor): Per-token CE losses from the model.
        model: The model (unused here, kept for API consistency with pretrain_gpt.py).

    Returns:
        Tuple of (loss scalar, num_tokens int, report dict).
    """
    losses = output_tensor.view(-1).float()
    loss_mask = loss_mask.view(-1).float()
    loss = torch.sum(losses * loss_mask)

    # num_tokens counts non-zero positions in the diffusion loss mask.
    num_tokens = (loss_mask > 0).sum().clone().detach().to(torch.int)
    report = {'diffusion loss': torch.cat([loss.clone().detach().view(1), num_tokens.view(1)])}

    # Check individual rank losses are not NaN/Inf prior to DP all-reduce.
    rerun_state_machine = get_rerun_state_machine()
    args = get_args()
    if args.check_for_nan_in_loss_and_grad:
        rerun_state_machine.validate_result(
            result=loss,
            rejection_func=torch.isnan,
            message="found NaN in local forward loss calculation",
            tolerance=0.0,
            fatal=True,
        )
        rerun_state_machine.validate_result(
            result=loss,
            rejection_func=torch.isinf,
            message="found Inf in local forward loss calculation",
            tolerance=0.0,
            fatal=True,
        )

    return loss, num_tokens, report


def diffusion_forward_step(data_iterator, model):
    """Forward step for masked diffusion language modeling.

    Implements the MDLM training objective from:
        Simple and Effective Masked Diffusion Language Models (https://arxiv.org/abs/2406.07524)

    The loss computation is designed to exactly match dllm's MDLMTrainer.compute_loss():
        token_nll = CE(logits, clean_input_ids) * loss_weights * masked_mask / norm_factor
    We encode loss_weights * masked_mask / norm_factor into diffusion_loss_mask so that
    Megatron's standard sum(CE * loss_mask) gives the identical scalar.

    Args:
        data_iterator: Input data iterator.
        model: The DiffusionLanguageModel.

    Returns:
        Tuple of (output_tensor, loss_func_partial).
    """
    args = get_args()
    timers = get_timers()

    # Get the batch.
    timers('batch-generator', log_level=2).start()
    global stimer
    with stimer(bdata=True):
        vp_stage = get_attr_wrapped_model(model, "vp_stage")
        tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = get_batch(
            data_iterator, vp_stage
        )
    timers('batch-generator').stop()

    with stimer:
        # Determine pipeline stage role.
        pp_size = args.pipeline_model_parallel_size
        is_pp_first = parallel_state.is_pipeline_first_stage()
        is_pp_last = parallel_state.is_pipeline_last_stage()

        if tokens is None and not is_pp_first and not is_pp_last:
            # Middle PP stage: no data, just pass through.
            output_tensor = model(None, None, None)
            return output_tensor, partial(loss_func, loss_mask, model=model)

        # --- Helper: compute diffusion_loss_mask from loss_mask using deterministic RNG ---
        def _compute_diffusion_loss_mask(loss_mask_local, b, l, device):
            """Compute diffusion_loss_mask deterministically from loss_mask.

            Uses a CPU RNG seeded by (iteration, dp_rank, cp_rank) so all TP ranks
            and PP stages produce identical masks for the same micro-batch.
            """
            scheduler = _get_scheduler()
            epsilon = getattr(args, 'diffusion_time_epsilon', 1e-3)

            dp_rank = parallel_state.get_data_parallel_rank()
            cp_rank = parallel_state.get_context_parallel_rank()
            iteration = getattr(args, 'curr_iteration', 0)
            cpu_rng = torch.Generator(device='cpu')
            cpu_rng.manual_seed(iteration * 100003 + dp_rank * 997 + cp_rank)

            t = epsilon + (1.0 - epsilon) * torch.rand(b, generator=cpu_rng)
            t_device = t.to(device=device, dtype=torch.float32)

            alpha_t = scheduler(t_device)
            p_mask = (1.0 - alpha_t).unsqueeze(1).expand(b, l)

            maskable_mask = (loss_mask_local > 0)
            rand_mask = torch.rand(b, l, generator=cpu_rng).to(device=device)
            masked_mask = (rand_mask < p_mask) & maskable_mask

            loss_weight_type = getattr(args, 'diffusion_loss_weight_type', 'scheduler')
            if loss_weight_type == 'scheduler':
                loss_weights = scheduler.weight(t_device).unsqueeze(1).repeat(1, l)
            else:
                loss_weights = torch.ones(b, l, device=device)

            diff_loss_mask = loss_weights * masked_mask.float()

            loss_norm_type = getattr(args, 'diffusion_loss_norm_type', 'token')
            if loss_norm_type == 'token':
                diff_loss_mask = diff_loss_mask / maskable_mask.sum().clamp(min=1)
            elif loss_norm_type == 'sequence':
                diff_loss_mask = diff_loss_mask / (
                    maskable_mask.sum(-1, keepdim=True).clamp(min=1) * b
                )
            elif loss_norm_type == 'batch':
                diff_loss_mask = diff_loss_mask / b
            else:
                raise ValueError(f"Invalid diffusion_loss_norm_type: {loss_norm_type}")

            return masked_mask, diff_loss_mask

        if is_pp_first and tokens is not None:
            # === First PP stage: apply diffusion masking to tokens ===
            b, l = tokens.shape
            device = tokens.device
            tp_rank = parallel_state.get_tensor_model_parallel_rank()
            tp_group = parallel_state.get_tensor_model_parallel_group()
            tp_src = parallel_state.get_tensor_model_parallel_src_rank()

            if tp_rank == 0:
                masked_mask, diffusion_loss_mask = _compute_diffusion_loss_mask(
                    loss_mask, b, l, device
                )
                mask_token_id = getattr(args, 'diffusion_mask_token_id', 128255)
                noised_input_ids = torch.where(
                    masked_mask, torch.tensor(mask_token_id, device=device), tokens
                )
            else:
                noised_input_ids = torch.empty_like(tokens)
                diffusion_loss_mask = torch.empty(b, l, dtype=torch.float32, device=device)

            torch.distributed.broadcast(noised_input_ids, tp_src, group=tp_group)
            torch.distributed.broadcast(diffusion_loss_mask, tp_src, group=tp_group)

            output_tensor = model(
                noised_input_ids, position_ids, attention_mask,
                labels=tokens, loss_mask=diffusion_loss_mask,
                packed_seq_params=packed_seq_params,
            )

        elif is_pp_last:
            # === Last PP stage: compute diffusion_loss_mask for loss, model gets hidden_states from pipeline ===
            tp_rank = parallel_state.get_tensor_model_parallel_rank()
            tp_group = parallel_state.get_tensor_model_parallel_group()
            tp_src = parallel_state.get_tensor_model_parallel_src_rank()

            if tp_rank == 0 and loss_mask is not None:
                b, l = loss_mask.shape
                device = loss_mask.device
                _, diffusion_loss_mask = _compute_diffusion_loss_mask(
                    loss_mask, b, l, device
                )
            else:
                # Get shape from labels (broadcast to all TP ranks on last stage)
                if labels is not None:
                    b, l = labels.shape
                    device = labels.device
                else:
                    # Fallback: use args
                    b = args.micro_batch_size
                    l = args.seq_length
                    device = torch.cuda.current_device()
                diffusion_loss_mask = torch.empty(b, l, dtype=torch.float32, device=device)

            torch.distributed.broadcast(diffusion_loss_mask, tp_src, group=tp_group)

            # Model receives hidden states from pipeline; tokens/position_ids not used.
            output_tensor = model(
                tokens, position_ids, attention_mask,
                labels=labels, loss_mask=diffusion_loss_mask,
                packed_seq_params=packed_seq_params,
            )

        else:
            # PP=1 case or unexpected: standard handling
            output_tensor = model(tokens, position_ids, attention_mask)
            diffusion_loss_mask = loss_mask

    return output_tensor, partial(loss_func, diffusion_loss_mask, model=model)


def is_dataset_built_on_rank(vp_stage=None):
    return (
        is_first_or_last_pipeline_stage(vp_stage)
        and parallel_state.get_tensor_model_parallel_rank() == 0
    )


def core_gpt_dataset_config_from_args(args):
    tokenizer = get_tokenizer()

    blend, blend_per_split = get_blend_and_blend_per_split(args)

    return GPTDatasetConfig(
        random_seed=args.seed,
        sequence_length=args.seq_length,
        blend=blend,
        blend_per_split=blend_per_split,
        split=args.split,
        num_dataset_builder_threads=args.num_dataset_builder_threads,
        path_to_cache=args.data_cache_path,
        mmap_bin_files=args.mmap_bin_files,
        tokenizer=tokenizer,
        reset_position_ids=args.reset_position_ids,
        reset_attention_mask=args.reset_attention_mask,
        eod_mask_loss=args.eod_mask_loss,
        create_attention_mask=args.create_attention_mask_in_dataloader,
        context_parallel_size=args.context_parallel_size,
        data_parallel_size=args.data_parallel_size,
    )


def train_valid_test_datasets_provider(train_val_test_num_samples):
    """Build the train, validation, and test datasets.

    Args:
        train_val_test_num_samples: A list containing the number of samples in
            train, validation, and test splits.
    """
    args = get_args()

    print_rank_0('> building train, validation, and test datasets for Diffusion LM ...')

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

    config = core_gpt_dataset_config_from_args(args)

    train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
        dataset_type,
        train_val_test_num_samples,
        is_dataset_built_on_rank,
        config,
    ).build()

    print_rank_0('> finished creating Diffusion LM datasets ...')

    return train_ds, valid_ds, test_ds


if __name__ == "__main__":
    # Timestamp right after entering __main__ block (after all imports/library setup).
    _MAIN_ENTRY_TIME = time.time()

    # Register startup timestamps for timing report in pretrain().
    set_startup_timestamps(program_start=_PROGRAM_START_TIME, main_entry=_MAIN_ENTRY_TIME)

    train_valid_test_datasets_provider.is_distributed = True

    pretrain(
        train_valid_test_datasets_provider,
        partial(model_provider, diffusion_builder),
        ModelType.encoder_or_decoder,
        diffusion_forward_step,
        args_defaults={'tokenizer_type': 'NullTokenizer'},
    )