diffusers/examples/discrete_diffusion/train_llada2.py

#!/usr/bin/env python
# Copyright 2025 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import math
import os
from dataclasses import asdict, dataclass
from typing import Dict, Optional

import torch
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration, set_seed
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer, DataCollatorForLanguageModeling, get_scheduler

from diffusers import BlockRefinementScheduler
from diffusers.training_utils import compute_confidence_aware_loss


logger = get_logger(__name__)


@dataclass
class TrainConfig:
    model_name_or_path: str
    dataset_name: str
    dataset_config_name: Optional[str]
    text_column: str
    cache_dir: Optional[str]
    use_dummy_data: bool
    num_dummy_samples: int

    output_dir: str
    seed: int
    max_train_steps: int
    checkpointing_steps: int
    logging_steps: int

    per_device_train_batch_size: int
    gradient_accumulation_steps: int
    learning_rate: float
    weight_decay: float
    lr_scheduler: str
    lr_warmup_steps: int

    max_length: int
    prompt_length: int
    block_length: int

    lambda_conf: float
    conf_temperature: float


def parse_args() -> TrainConfig:
    parser = argparse.ArgumentParser(description="Train block-refinement with a confidence-aware loss on a causal LM.")

    parser.add_argument("--model_name_or_path", type=str, default="Qwen/Qwen2.5-0.5B")
    parser.add_argument("--dataset_name", type=str, default="wikitext")
    parser.add_argument("--dataset_config_name", type=str, default="wikitext-2-raw-v1")
    parser.add_argument("--text_column", type=str, default="text")
    parser.add_argument("--cache_dir", type=str, default=None)
    parser.add_argument("--use_dummy_data", action="store_true", help="Use random-token data instead of downloading.")
    parser.add_argument("--num_dummy_samples", type=int, default=2048)

    parser.add_argument("--output_dir", type=str, default="block-refinement-output")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--max_train_steps", type=int, default=1000)
    parser.add_argument("--checkpointing_steps", type=int, default=500)
    parser.add_argument("--logging_steps", type=int, default=50)

    parser.add_argument("--per_device_train_batch_size", type=int, default=1)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=8)
    parser.add_argument("--learning_rate", type=float, default=2e-5)
    parser.add_argument("--weight_decay", type=float, default=0.0)
    parser.add_argument(
        "--lr_scheduler", type=str, default="cosine", choices=["linear", "cosine", "cosine_with_restarts"]
    )
    parser.add_argument("--lr_warmup_steps", type=int, default=100)

    parser.add_argument("--max_length", type=int, default=256)
    parser.add_argument("--prompt_length", type=int, default=32)
    parser.add_argument("--block_length", type=int, default=32)

    parser.add_argument("--lambda_conf", type=float, default=2.0)
    parser.add_argument("--conf_temperature", type=float, default=0.5)

    args = parser.parse_args()
    return TrainConfig(**vars(args))


def tokenize_fn(examples: Dict, tokenizer, text_column: str, max_length: int):
    texts = examples[text_column]
    texts = [t for t in texts if isinstance(t, str) and len(t.strip()) > 0]
    return tokenizer(texts, truncation=True, padding=False, max_length=max_length)


class RandomTokenDataset(torch.utils.data.Dataset):
    def __init__(self, *, num_samples: int, seq_len: int, vocab_size: int, pad_token_id: int):
        self.num_samples = int(num_samples)
        self.seq_len = int(seq_len)
        self.vocab_size = int(vocab_size)
        self.pad_token_id = int(pad_token_id)

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        del idx
        input_ids = torch.randint(0, self.vocab_size, (self.seq_len,), dtype=torch.long)
        attention_mask = torch.ones_like(input_ids)
        return {"input_ids": input_ids, "attention_mask": attention_mask}


def main():
    cfg = parse_args()
    if cfg.prompt_length >= cfg.max_length:
        raise ValueError("`prompt_length` must be < `max_length`.")
    if cfg.block_length <= 0:
        raise ValueError("`block_length` must be > 0.")

    project_config = ProjectConfiguration(project_dir=cfg.output_dir, logging_dir=os.path.join(cfg.output_dir, "logs"))
    accelerator = Accelerator(
        gradient_accumulation_steps=cfg.gradient_accumulation_steps,
        project_config=project_config,
    )
    if accelerator.is_main_process:
        os.makedirs(cfg.output_dir, exist_ok=True)
    accelerator.wait_for_everyone()

    set_seed(cfg.seed)
    logger.info("Training configuration: %s", asdict(cfg))

    tokenizer = AutoTokenizer.from_pretrained(cfg.model_name_or_path, use_fast=True, cache_dir=cfg.cache_dir)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    if tokenizer.mask_token_id is None:
        tokenizer.add_special_tokens({"mask_token": "[MASK]"})

    load_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
    model = AutoModelForCausalLM.from_pretrained(cfg.model_name_or_path, cache_dir=cfg.cache_dir, dtype=load_dtype)
    model.resize_token_embeddings(len(tokenizer))
    if load_dtype == torch.float32:
        model.to(dtype=torch.float32)

    mask_token_id = int(tokenizer.mask_token_id)

    if cfg.use_dummy_data:
        dataset = RandomTokenDataset(
            num_samples=cfg.num_dummy_samples,
            seq_len=cfg.max_length,
            vocab_size=len(tokenizer),
            pad_token_id=int(tokenizer.pad_token_id),
        )
        train_dataloader = DataLoader(
            dataset,
            shuffle=True,
            batch_size=cfg.per_device_train_batch_size,
            drop_last=True,
        )
    else:
        raw_datasets = load_dataset(cfg.dataset_name, cfg.dataset_config_name, cache_dir=cfg.cache_dir)
        if "train" not in raw_datasets:
            raise ValueError(f"Dataset {cfg.dataset_name} has no 'train' split.")

        with accelerator.main_process_first():
            tokenized = raw_datasets["train"].map(
                lambda ex: tokenize_fn(ex, tokenizer, cfg.text_column, cfg.max_length),
                batched=True,
                remove_columns=raw_datasets["train"].column_names,
                desc="Tokenizing",
            )

        collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False, return_tensors="pt")
        train_dataloader = DataLoader(
            tokenized, shuffle=True, collate_fn=collator, batch_size=cfg.per_device_train_batch_size, drop_last=True
        )

    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / cfg.gradient_accumulation_steps)
    num_train_epochs = math.ceil(cfg.max_train_steps / num_update_steps_per_epoch)

    lr_scheduler = get_scheduler(
        name=cfg.lr_scheduler,
        optimizer=optimizer,
        num_warmup_steps=cfg.lr_warmup_steps,
        num_training_steps=cfg.max_train_steps,
    )

    model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, lr_scheduler
    )

    noise_scheduler = BlockRefinementScheduler(block_length=cfg.block_length)

    global_step = 0
    model.train()

    for _epoch in range(num_train_epochs):
        for batch in train_dataloader:
            with accelerator.accumulate(model):
                input_ids = batch["input_ids"]
                attention_mask = batch.get("attention_mask", torch.ones_like(input_ids))

                gen = torch.Generator(device=input_ids.device).manual_seed(cfg.seed + global_step)
                noisy, noisy_rev, masked, masked_rev = noise_scheduler.add_noise(
                    input_ids,
                    attention_mask,
                    prompt_length=cfg.prompt_length,
                    block_length=cfg.block_length,
                    mask_token_id=mask_token_id,
                    generator=gen,
                )

                position_ids = (
                    torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0).expand_as(input_ids)
                )

                logits = model(input_ids=noisy, attention_mask=attention_mask, position_ids=position_ids).logits
                logits_rev = model(
                    input_ids=noisy_rev, attention_mask=attention_mask, position_ids=position_ids
                ).logits

                logits = logits.clone()
                logits[..., mask_token_id] = torch.finfo(logits.dtype).min
                logits_rev = logits_rev.clone()
                logits_rev[..., mask_token_id] = torch.finfo(logits_rev.dtype).min

                valid = attention_mask.to(dtype=torch.bool)
                masked = masked & valid
                masked_rev = masked_rev & valid

                labels = input_ids.clone()
                labels[~masked] = -100
                labels_rev = input_ids.clone()
                labels_rev[~masked_rev] = -100

                weights = masked.to(dtype=logits.dtype)
                weights_rev = masked_rev.to(dtype=logits.dtype)

                loss, loss_sft, loss_conf = compute_confidence_aware_loss(
                    logits,
                    labels,
                    lambda_conf=cfg.lambda_conf,
                    temperature=cfg.conf_temperature,
                    per_token_weights=weights,
                )
                loss_rev, loss_sft_rev, loss_conf_rev = compute_confidence_aware_loss(
                    logits_rev,
                    labels_rev,
                    lambda_conf=cfg.lambda_conf,
                    temperature=cfg.conf_temperature,
                    per_token_weights=weights_rev,
                )

                total_loss = loss + loss_rev
                accelerator.backward(total_loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad(set_to_none=True)

            if accelerator.sync_gradients:
                global_step += 1

                if global_step % cfg.logging_steps == 0 and accelerator.is_main_process:
                    logger.info(
                        "step=%d loss=%.4f sft=%.4f conf=%.4f lr=%.6g",
                        global_step,
                        total_loss.item(),
                        (loss_sft + loss_sft_rev).item(),
                        (loss_conf + loss_conf_rev).item(),
                        lr_scheduler.get_last_lr()[0],
                    )
                    print(
                        f"step={global_step} loss={total_loss.item():.4f} "
                        f"sft={(loss_sft + loss_sft_rev).item():.4f} "
                        f"conf={(loss_conf + loss_conf_rev).item():.4f} "
                        f"lr={lr_scheduler.get_last_lr()[0]:.6g}"
                    )

                if cfg.checkpointing_steps > 0 and global_step % cfg.checkpointing_steps == 0:
                    accelerator.wait_for_everyone()
                    if accelerator.is_main_process:
                        save_dir = os.path.join(cfg.output_dir, f"checkpoint-{global_step}")
                        os.makedirs(save_dir, exist_ok=True)
                        accelerator.unwrap_model(model).save_pretrained(save_dir, save_function=accelerator.save)
                        tokenizer.save_pretrained(save_dir)

                if global_step >= cfg.max_train_steps:
                    break

        if global_step >= cfg.max_train_steps:
            break

    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        final_dir = os.path.join(cfg.output_dir, "final")
        os.makedirs(final_dir, exist_ok=True)
        accelerator.unwrap_model(model).save_pretrained(final_dir, save_function=accelerator.save)
        tokenizer.save_pretrained(final_dir)

    logger.info("Done.")


if __name__ == "__main__":
    main()