train.py

# train_pp.py
# Run with: uv sync && uv run torchrun --nproc_per_node=2 experiments/pp-ft/train.py
#
# Pipeline Parallel + LoRA finetuning of Qwen3-4B on AMI meeting transcripts.
# Uses PyTorch's native torch.distributed.pipelining (PipelineStage + Schedule1F1B).
#
# --- How PP differs from the other strategies in this repo ---
#
# DDP/FSDP/TP all replicate or shard the *same* layers across ranks. PP splits
# the model along its *depth*: each rank holds a contiguous slice of decoder
# layers and only those layers, and the activations stream rank-to-rank.
#
# For Qwen3-4B with 36 layers and PP=2:
#   - rank 0 holds: embed_tokens + layers[0:18]
#   - rank 1 holds: layers[18:36] + final norm + lm_head
#
# Each rank only stores ~half the parameters. The forward pass produces a
# hidden state on rank 0, sends it to rank 1, which produces logits and the
# loss. The backward pass flows in reverse: gradient on the loss → rank 1
# computes gradients for its layers and sends the input gradient back to rank
# 0 → rank 0 computes gradients for its layers.
#
# --- Why microbatching ---
#
# The naive PP forward leaves rank 1 idle while rank 0 is computing, and vice
# versa during backward. This idle time is the "pipeline bubble." Microbatching
# splits each batch into N microbatches and overlaps them: while rank 1 is
# processing microbatch i, rank 0 starts microbatch i+1. The bubble shrinks
# from 50% (no microbatching) to (num_stages - 1) / (num_stages - 1 + n_micro).
# With 2 stages and 8 microbatches: bubble = 1/9 ≈ 11%.
#
# Schedule1F1B (one-forward-one-backward) is the standard production schedule.
# Each rank alternates F and B passes once the pipeline is filled, which keeps
# the activation memory bounded to ~2 microbatches per rank instead of all of
# them (the GPipe behavior).
#
# --- LoRA + PP ---
#
# Unlike TP, PP does not replace nn.Linear with anything. PEFT's LoraLinear
# wrappers work without modification. The trick is that LoRA is applied to the
# *full* model first; then on each rank we keep only the slice of layers
# assigned to that stage. Each rank's optimizer sees only its local LoRA
# adapters. Saving requires each rank to write its own adapter shard, just
# like the TP scripts.
#
# --- Caveats ---
#
# PP at this scale is a teaching demo, not a production fit. Qwen3-4B fits
# comfortably on one 24 GB GPU, so PP is paying the bubble cost for a workload
# that did not need PP. The interesting numbers are the bubble percentage and
# the throughput vs DDP+LoRA at the same effective batch. PP shines at model
# sizes where a single GPU genuinely cannot hold the full model (e.g. 70B+
# without quantization), where it composes with TP and DP to form 3D
# parallelism in production stacks.

import csv
import os
import time
from typing import Optional

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from datasets import load_dataset
from peft import LoraConfig, TaskType, get_peft_model
from torch.distributed.pipelining import PipelineStage, Schedule1F1B
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader, Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")


# --- config ---

MODEL_ID   = "Qwen/Qwen3-4B"
OUTPUT_DIR = "checkpoints/pp"
LOG_FILE   = "logs/pp.csv"

DATASET = "ami"

MAX_LENGTH      = 512
BATCH_SIZE      = 8     # samples per schedule.step(); split across microbatches
N_MICROBATCHES  = 8     # microbatch_size = BATCH_SIZE / N_MICROBATCHES = 1
GRAD_ACCUM      = 2     # 2 schedule.steps before optimizer.step
                        # effective batch = BATCH_SIZE * GRAD_ACCUM = 16
EPOCHS          = 3
LR              = 2e-4
MAX_GRAD_NORM   = 1.0


# --- distributed setup ---

def setup() -> tuple[int, int, int]:
    dist.init_process_group(backend="nccl")
    rank        = dist.get_rank()
    world_size  = dist.get_world_size()
    local_rank  = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return rank, world_size, local_rank


def cleanup():
    dist.destroy_process_group()


# --- prompt formatting ---

SYSTEM_PROMPT = (
    "You are a meeting assistant. Given a meeting transcript, write a concise summary "
    "that covers the key discussion points, decisions made, and action items with "
    "owners and deadlines where mentioned."
)


# --- pipeline stage wrapper ---
#
# PyTorch's PipelineStage wraps an nn.Module whose forward signature is
# stage-aware: stage 0 takes the raw input (input_ids), intermediate and final
# stages take a hidden state tensor, and the final stage returns logits.
#
# We don't subclass Qwen3ForCausalLM. Instead we mirror what Qwen3Model.forward
# does internally (embed → for layer in layers → norm) but only over our
# stage's slice of layers. This is fragile across transformers versions
# because Qwen3DecoderLayer.forward keeps gaining new kwargs. The currently
# expected calling convention is:
#
#   layer(hidden_states,
#         attention_mask=None,
#         position_ids=position_ids,
#         past_key_value=None,
#         output_attentions=False,
#         use_cache=False,
#         cache_position=cache_position,
#         position_embeddings=(cos, sin))
#
# `position_embeddings` was introduced in transformers 4.45 and is now required;
# we precompute (cos, sin) once via the model's rotary_emb module and pass them
# into each layer.

class PPQwen3Stage(nn.Module):
    def __init__(self, peft_model, stage_idx: int, num_stages: int):
        super().__init__()
        self.stage_idx  = stage_idx
        self.num_stages = num_stages

        # peft_model is PeftModel(Qwen3ForCausalLM). Reach inside to grab the
        # layers and the auxiliary modules.
        # Structure: peft_model.base_model.model is the underlying
        # Qwen3ForCausalLM after PEFT wrapping. .model is the Qwen3Model.
        causal_lm = peft_model.base_model.model
        inner     = causal_lm.model

        all_layers       = list(inner.layers)
        n_layers         = len(all_layers)
        layers_per_stage = n_layers // num_stages
        start = stage_idx * layers_per_stage
        end   = start + layers_per_stage if stage_idx < num_stages - 1 else n_layers
        self.layers = nn.ModuleList(all_layers[start:end])

        if stage_idx == 0:
            self.embed_tokens = inner.embed_tokens
        else:
            self.embed_tokens = None

        if stage_idx == num_stages - 1:
            self.norm    = inner.norm
            self.lm_head = causal_lm.lm_head
        else:
            self.norm    = None
            self.lm_head = None

        # Rotary embeddings live on the parent Qwen3Model. We hold a reference
        # so every stage can compute (cos, sin) for its slice independently.
        self.rotary_emb = inner.rotary_emb

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Stage 0 receives input_ids [b, s]; subsequent stages receive a
        # hidden state tensor [b, s, hidden] from the previous stage.
        if self.stage_idx == 0:
            hidden = self.embed_tokens(x)
        else:
            hidden = x

        b, s = hidden.shape[:2]
        device = hidden.device

        # Position ids and cache_position are independent of the stage; both
        # ranks see the same sequence indices.
        position_ids   = torch.arange(s, device=device).unsqueeze(0).expand(b, -1)
        cache_position = torch.arange(s, device=device)

        # rotary_emb takes the hidden state (for dtype/device inference) and
        # position_ids, returns the (cos, sin) tuple expected by attention.
        position_embeddings = self.rotary_emb(hidden, position_ids)

        for layer in self.layers:
            layer_out = layer(
                hidden,
                attention_mask=None,           # SDPA infers causal mask
                position_ids=position_ids,
                past_key_value=None,
                output_attentions=False,
                use_cache=False,
                cache_position=cache_position,
                position_embeddings=position_embeddings,
            )
            hidden = layer_out[0] if isinstance(layer_out, tuple) else layer_out

        if self.stage_idx == self.num_stages - 1:
            hidden = self.norm(hidden)
            logits = self.lm_head(hidden)
            return logits
        return hidden


# --- model build ---

def build_model_and_stage(rank: int, num_stages: int, local_rank: int) -> tuple[PPQwen3Stage, AutoTokenizer]:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    if rank == 0:
        print("loading HF Qwen3-4B in BF16 on CPU...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.bfloat16,
        token=HF_TOKEN,
        low_cpu_mem_usage=True,
    )
    model.config.use_cache = False

    # Apply LoRA to the full model on every rank. PEFT wraps the target
    # nn.Linears in-place, so when we slice out a stage's layers below the
    # adapters come along automatically.
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=16,
        lora_alpha=32,
        lora_dropout=0.15,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
        bias="none",
    )
    peft_model = get_peft_model(model, lora_config)

    # Build this rank's stage. Move it to GPU. The full HF model is still in
    # CPU memory at this point; we drop the reference after slicing so the
    # unused half can be garbage-collected.
    stage_module = PPQwen3Stage(peft_model, stage_idx=rank, num_stages=num_stages)
    stage_module = stage_module.to(f"cuda:{local_rank}")

    # Drop the full peft_model — the layers we kept are now owned by stage_module.
    del peft_model
    del model

    if rank == 0:
        n_train = sum(p.numel() for p in stage_module.parameters() if p.requires_grad)
        n_total = sum(p.numel() for p in stage_module.parameters())
        print(f"stage 0 trainable params: {n_train:,} / {n_total:,} ({100*n_train/n_total:.2f}%)")

    return stage_module, tokenizer


# --- dataset ---

def load_ami() -> list[dict]:
    train_split = load_dataset("knkarthick/AMI", split="train",      token=HF_TOKEN)
    val_split   = load_dataset("knkarthick/AMI", split="validation", token=HF_TOKEN)
    samples = []
    for row in list(train_split) + list(val_split):
        t, s = row["dialogue"].strip(), row["summary"].strip()
        if t and s:
            samples.append({"transcript": t, "summary": s})
    return samples


class MeetingDataset(Dataset):
    def __init__(self, samples: list[dict], tokenizer, max_length: int = MAX_LENGTH):
        self.samples    = samples
        self.tokenizer  = tokenizer
        self.max_length = max_length

    def __len__(self): return len(self.samples)

    def __getitem__(self, idx):
        s = self.samples[idx]
        prompt = (
            f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
            f"<|im_start|>user\n{s['transcript']}<|im_end|>\n"
            f"<|im_start|>assistant\n"
        )
        completion = s["summary"] + "<|im_end|>"
        prompt_ids     = self.tokenizer(prompt,     add_special_tokens=False, max_length=self.max_length, truncation=True)["input_ids"]
        completion_ids = self.tokenizer(completion, add_special_tokens=False, max_length=self.max_length, truncation=True)["input_ids"]
        max_prompt_len = self.max_length - len(completion_ids)
        prompt_ids = prompt_ids[-max_prompt_len:]
        ids = torch.tensor(prompt_ids + completion_ids, dtype=torch.long)
        labels = ids.clone()
        labels[: len(prompt_ids)] = -100
        return ids, labels


def collate_fn(batch):
    ids, lbls = zip(*batch)
    max_len = max(x.size(0) for x in ids)
    pad_ids = torch.zeros(len(ids),  max_len, dtype=torch.long)
    pad_lbl = torch.full((len(lbls), max_len), fill_value=-100, dtype=torch.long)
    for i, (a, b) in enumerate(zip(ids, lbls)):
        pad_ids[i, : a.size(0)] = a
        pad_lbl[i, : b.size(0)] = b
    return pad_ids, pad_lbl


# --- save ---

def save_adapter(stage_module: PPQwen3Stage, tokenizer, save_path: str, rank: int):
    """Each rank saves its own LoRA adapter shard. The full adapter is the
    union of all rank shards (no overlap, since each layer lives on exactly
    one stage)."""
    adapter_state = {n: p.detach().cpu() for n, p in stage_module.named_parameters() if p.requires_grad}
    rank_path = f"{save_path}/rank_{rank}"
    os.makedirs(rank_path, exist_ok=True)
    torch.save(adapter_state, f"{rank_path}/adapter_model.bin")
    if rank == 0:
        tokenizer.save_pretrained(save_path)
        print(f"saved adapter shards -> {save_path}/rank_*/  ({len(adapter_state)} tensors on rank 0)")
    dist.barrier()


# --- metrics logging ---

LOG_COLUMNS = ["strategy", "epoch", "step", "loss", "tokens_per_sec",
               "step_time_sec", "peak_vram_gb", "lr"]


def setup_logger(path: str) -> csv.DictWriter:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    f = open(path, "w", newline="")
    writer = csv.DictWriter(f, fieldnames=LOG_COLUMNS)
    writer.writeheader()
    return writer


# --- training loop ---

def loss_fn(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
    """Causal LM cross-entropy with shifted labels. -100 is ignored."""
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()
    return F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)).float(),
        shift_labels.view(-1),
        ignore_index=-100,
    )


def train(rank: int, world_size: int, local_rank: int):
    stage_module, tokenizer = build_model_and_stage(rank, world_size, local_rank)

    # Build the PipelineStage. PyTorch needs the module, this rank's stage
    # index, the total number of stages, and the device.
    stage = PipelineStage(
        stage_module,
        stage_index=rank,
        num_stages=world_size,
        device=torch.device(f"cuda:{local_rank}"),
    )

    # Schedule1F1B (one-forward-one-backward): once the pipeline is filled, each
    # rank alternates a forward microbatch and a backward microbatch. Activation
    # memory is bounded to ~2 microbatches per rank instead of all of them.
    schedule = Schedule1F1B(stage, n_microbatches=N_MICROBATCHES, loss_fn=loss_fn)

    samples = load_ami()
    dataset = MeetingDataset(samples, tokenizer)

    # All ranks see the same data in the same order (PP is not data-parallel).
    g = torch.Generator()
    g.manual_seed(42)
    loader = DataLoader(
        dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        collate_fn=collate_fn,
        pin_memory=True,
        num_workers=0,
        generator=g,
        drop_last=True,            # PP needs every batch to have BATCH_SIZE samples
                                    # so the microbatch split is even
    )

    optimizer = AdamW(
        filter(lambda p: p.requires_grad, stage_module.parameters()),
        lr=LR, betas=(0.9, 0.999), weight_decay=0.01,
    )
    total_update_steps = (len(loader) // GRAD_ACCUM) * EPOCHS
    scheduler = CosineAnnealingLR(optimizer, T_max=total_update_steps, eta_min=LR * 0.1)

    logger = setup_logger(LOG_FILE) if rank == 0 else None

    stage_module.train()
    global_step = 0
    tokens_this_window = 0
    window_start       = time.perf_counter()

    for epoch in range(EPOCHS):
        g.manual_seed(42 + epoch)
        optimizer.zero_grad(set_to_none=True)

        for step, (input_ids, labels) in enumerate(loader):
            input_ids = input_ids.to(local_rank, non_blocking=True)
            labels    = labels.to(local_rank, non_blocking=True)
            tokens_this_window += input_ids.numel()

            # Drive the schedule. Stage 0 passes the input; the last stage
            # passes the target and a list to capture per-microbatch losses.
            losses_buf: list[torch.Tensor] = []
            if rank == 0:
                schedule.step(input_ids)
            elif rank == world_size - 1:
                schedule.step(target=labels, losses=losses_buf)
            else:
                schedule.step()

            # The last stage is the only rank that computed the loss. Broadcast
            # it to rank 0 so the logger can write a real number instead of 0.
            loss_tensor = torch.zeros(1, device=f"cuda:{local_rank}", dtype=torch.float32)
            if rank == world_size - 1 and losses_buf:
                stacked = torch.stack([l.detach().float() for l in losses_buf])
                loss_tensor[0] = stacked.mean()
            dist.broadcast(loss_tensor, src=world_size - 1)
            local_loss_val = loss_tensor.item()

            if (step + 1) % GRAD_ACCUM == 0:
                torch.nn.utils.clip_grad_norm_(
                    [p for p in stage_module.parameters() if p.requires_grad],
                    MAX_GRAD_NORM,
                )
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad(set_to_none=True)
                global_step += 1

                if rank == 0:
                    step_time      = time.perf_counter() - window_start
                    # PP: every rank sees the same tokens (data-replicated, not
                    # data-parallel), so rank 0's count IS the aggregate work.
                    tokens_per_sec = tokens_this_window / step_time
                    peak_vram_gb   = torch.cuda.memory_allocated(local_rank) / 1e9
                    lr             = scheduler.get_last_lr()[0]
                    print(
                        f"epoch {epoch} | step {global_step:4d} | loss {local_loss_val:.4f} "
                        f"| {tokens_per_sec:,.0f} tok/s | {step_time:.1f}s | "
                        f"{peak_vram_gb:.1f} GB VRAM"
                    )
                    logger.writerow({
                        "strategy":       "pp",
                        "epoch":          epoch,
                        "step":           global_step,
                        "loss":           round(local_loss_val, 4),
                        "tokens_per_sec": round(tokens_per_sec, 1),
                        "step_time_sec":  round(step_time, 2),
                        "peak_vram_gb":   round(peak_vram_gb, 2),
                        "lr":             f"{lr:.2e}",
                    })
                    tokens_this_window = 0
                    window_start       = time.perf_counter()

        save_adapter(stage_module, tokenizer, f"{OUTPUT_DIR}/epoch_{epoch}", rank)


def main():
    rank, world_size, local_rank = setup()
    try:
        train(rank, world_size, local_rank)
    finally:
        cleanup()


if __name__ == "__main__":
    main()