# ML Coding

## Data loading

### SFT / pre-training data loading

This is the data path used by the training skeleton below: **raw chat row → token IDs and prompt-masked labels → dynamic padding → shuffled mini-batches**. The worked example uses an SFT conversation. For plain pre-training, the loader has the same overall shape, but every real token is a target, so `labels` is a copy of `input_ids` instead of masking a prompt.

```python
import torch
import torch.nn.functional as F
from datasets import load_dataset
from torch.utils.data import DataLoader


def create_sft_dataloader(cfg, tokenizer):
    # Download/read the configured split as a Hugging Face Dataset. At this
    # point, each row still contains human-readable conversation dictionaries.
    raw = load_dataset(cfg.dataset_name, split=cfg.dataset_split)

    # One raw SFT example looks like this before any tokenization:
    #
    # raw[250]["messages"] = [
    #     {
    #         "role": "user",
    #         "content": "Can you write a haiku about two space cowboys "
    #                    "fighting against each other?",
    #     },
    #     {
    #         "role": "assistant",
    #         "content": "Lasers speeding through\n"
    #                    "A pair of space cowboys fight\n"
    #                    "The stars shine above",
    #     },
    # ]
    #
    # The row is ordinary Python data. No tensors exist yet.
    if cfg.max_samples is not None:
        raw = raw.select(range(min(cfg.max_samples, len(raw))))

    # Dataset.map sends a dictionary of columns, so batch["messages"] is a
    # list of conversations. This preprocessing batch size is unrelated to
    # the training batch size used by DataLoader later.
    def encode_batch(batch):
        # Keep the final assistant turn as the target and mask the prompt.
        conversations = [
            messages for messages in batch["messages"]
            if messages and messages[-1]["role"] == "assistant"
        ]
        if not conversations:
            return {"input_ids": [], "labels": []}

        # Render only the context and append the assistant-start marker.
        # Its length tells us exactly how much of the full sequence is prompt.
        prompt_ids = tokenizer.apply_chat_template(
            [messages[:-1] for messages in conversations],
            tokenize=True,
            add_generation_prompt=True,
            truncation=True,
            max_length=cfg.max_length,
        )

        # Render context + the expected final assistant answer. These IDs are
        # the sequence that the model reads during teacher-forced training.
        full_ids = tokenizer.apply_chat_template(
            conversations,
            tokenize=True,
            add_generation_prompt=False,
            truncation=True,
            max_length=cfg.max_length,
        )

        input_ids, labels = [], []
        for prompt, full in zip(prompt_ids, full_ids, strict=True):
            prompt_length = min(len(prompt), len(full))
            if prompt_length == len(full):
                continue  # Truncation removed every assistant target token.

            input_ids.append(full)
            labels.append([-100] * prompt_length + full[prompt_length:])

            # After encoding the haiku row with Qwen2.5-Instruct:
            #
            # len(full) = 64
            # prompt_length = 45
            # input_ids = [151644, 8948, 198, ..., 3403, 151645, 198]
            # labels    = [-100] * 45 +
            #             [43803, 388, 57575, ..., 3403, 151645, 198]
            #
            # input_ids contains the complete conversation. labels has the
            # same length, but its 45 prompt positions are -100, so loss is
            # computed only on the 19 assistant-answer positions.

        return {"input_ids": input_ids, "labels": labels}

    encoded = raw.map(
        encode_batch,
        batched=True,
        batch_size=256,
        remove_columns=raw.column_names,
        desc="Encoding conversations",
    ).with_format("torch")

    if len(encoded) == 0:
        raise RuntimeError("No trainable rows after tokenization.")

    # Fetching encoded[0] now returns two variable-length 1-D tensors:
    #
    # encoded[0]["input_ids"].shape == [sequence_length]
    # encoded[0]["labels"].shape    == [sequence_length]
    #
    # They remain on CPU. Only a completed mini-batch moves to the device.

    # Collator and DataLoader: dynamically right-pad each shuffled mini-batch.
    def collate(examples):
        max_len = max(example["input_ids"].size(0) for example in examples)
        input_ids, attention_mask, labels = [], [], []

        for example in examples:
            ids, target = example["input_ids"], example["labels"]
            pad = max_len - ids.size(0)
            input_ids.append(F.pad(ids, (0, pad), value=tokenizer.pad_token_id))
            attention_mask.append(F.pad(torch.ones_like(ids), (0, pad), value=0))
            labels.append(F.pad(target, (0, pad), value=-100))

        batch = {
            "input_ids": torch.stack(input_ids),
            "attention_mask": torch.stack(attention_mask),
            "labels": torch.stack(labels),
        }

        # Example: collating lengths 64 and 51 pads both rows to 64:
        #
        # batch["input_ids"].shape      == [2, 64]
        # batch["attention_mask"].shape == [2, 64]
        # batch["labels"].shape         == [2, 64]
        #
        # row 0 attention_mask: [1] * 64
        # row 1 attention_mask: [1] * 51 + [0] * 13
        # row 1 labels:         [... real targets ...] + [-100] * 13
        #
        # Padding IDs make a rectangular tensor; attention 0 hides padding
        # from the transformer, and label -100 hides it from cross-entropy.
        return batch

    dataloader = DataLoader(
        encoded,
        batch_size=cfg.batch_size,
        shuffle=True,
        collate_fn=collate,
        num_workers=0,
    )

    # Iteration now yields one training mini-batch at a time:
    # for batch in dataloader:
    #     batch = {name: tensor.to(device) for name, tensor in batch.items()}
    return dataloader
```

### Preference reward-model data loading

This path uses [`argilla/ultrafeedback-binarized-preferences-cleaned`](https://huggingface.co/datasets/argilla/ultrafeedback-binarized-preferences-cleaned): **raw preference pair → two tokenized conversations → filtering → train/validation split → separately padded chosen/rejected batches**.

```python
import os
import random
from typing import Dict, List

import torch
from datasets import Dataset, load_dataset, load_from_disk
from torch.utils.data import DataLoader
from transformers import AutoTokenizer


def format_conversation(messages: List[Dict]) -> str:
    """Format a conversation as a simple string."""
    result = []
    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        result.append(f"{role}: {content}")
    return "\n".join(result)


def tokenize_messages(
    tokenizer: AutoTokenizer,
    messages: List[Dict],
    max_length: int,
    return_tensors: str | None = None,
) -> Dict:
    """Tokenize messages with the tokenizer chat template when available."""
    if tokenizer.chat_template is not None:
        return tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=False,
            max_length=max_length,
            truncation=True,
            return_dict=True,
            return_tensors=return_tensors,
        )

    return tokenizer(
        format_conversation(messages),
        max_length=max_length,
        truncation=True,
        add_special_tokens=True,
        return_tensors=return_tensors,
    )


def build_preference_dataset(
    tokenizer: AutoTokenizer,
    config: Config,
) -> Dataset:
    """Build preference dataset from UltraFeedback.

    Each example contains:
    - chosen_ids: Token IDs for the chosen response
    - rejected_ids: Token IDs for the rejected response

    Pairs that are token-identical after truncation carry no training signal
    and are dropped, so the returned dataset can be smaller than `limit`.
    """
    random.seed(config.seed)

    dataset_name = config.dataset_name
    dataset_split = config.dataset_split
    if os.path.exists(dataset_name):
        raw = load_from_disk(dataset_name)
        if hasattr(raw, "keys"):
            raw = raw[dataset_split]
    else:
        raw = load_dataset(dataset_name, split=dataset_split)

    # For this example:
    # config.dataset_name =
    #     "argilla/ultrafeedback-binarized-preferences-cleaned"
    #
    # One raw row is a preference record with the same user request in both
    # conversations and two different assistant answers:
    #
    # raw[0] = {
    #     "prompt": "Can you write a C++ program that ... checks if a country "
    #               "borders the Mediterranean Sea?",
    #     "chosen": [
    #         {"role": "user", "content": "Can you write a C++ program ...?"},
    #         {"role": "assistant", "content": "Here's a C++ program ..."},
    #     ],
    #     "rejected": [
    #         {"role": "user", "content": "Can you write a C++ program ...?"},
    #         {"role": "assistant", "content": "Sure, here is the program ..."},
    #     ],
    #     # The dataset also has ratings/model metadata; this loader removes it.
    # }

    # Shuffle and limit
    raw = raw.shuffle(seed=config.seed).select(range(min(config.samples, len(raw))))

    records = []
    max_length = config.max_length
    dropped_identical = 0
    for ex in raw:
        # Extract prompt and responses
        prompt = ex.get("prompt", "")
        chosen = ex.get("chosen", [])
        rejected = ex.get("rejected", [])

        # Handle different dataset formats
        if isinstance(chosen, list) and len(chosen) > 0:
            # Conversation format
            chosen_messages = chosen
            rejected_messages = rejected
        elif isinstance(chosen, str):
            chosen_messages = [
                {"role": "user", "content": prompt},
                {"role": "assistant", "content": chosen},
            ]
            rejected_messages = [
                {"role": "user", "content": prompt},
                {"role": "assistant", "content": rejected},
            ]
        else:
            continue

        chosen_tokens = tokenize_messages(tokenizer, chosen_messages, max_length)
        rejected_tokens = tokenize_messages(tokenizer, rejected_messages, max_length)

        # The raw strings have now become two independent integer sequences.
        # Exact IDs depend on the tokenizer; structurally, one row is now:
        #
        # chosen_tokens = {
        #     "input_ids":     [BOS, user..., assistant..., EOS],
        #     "attention_mask": [  1,       1...,         1...,   1],
        # }
        # rejected_tokens = {
        #     "input_ids":     [BOS, user..., different_assistant..., EOS],
        #     "attention_mask": [  1,       1...,                   1...,   1],
        # }
        #
        # Their prompt prefixes normally match, but response tokens and total
        # lengths can differ. No reward-model labels are stored: by convention,
        # the Bradley–Terry target says chosen should score above rejected.

        # A long shared prompt can consume the whole max_length window, truncating
        # both responses away and leaving identical token ids. The Bradley-Terry
        # loss is a constant -log(0.5) with zero gradient on such pairs, so drop
        # them rather than left-truncate (which would cut the prompt the reward
        # head needs). See https://github.com/natolambert/rlhf-book/issues/503.
        if chosen_tokens["input_ids"] == rejected_tokens["input_ids"]:
            dropped_identical += 1
            continue

        records.append(
            {
                "chosen_ids": chosen_tokens["input_ids"],
                "chosen_mask": chosen_tokens["attention_mask"],
                "rejected_ids": rejected_tokens["input_ids"],
                "rejected_mask": rejected_tokens["attention_mask"],
            }
        )

        # The retained example has four variable-length Python lists:
        # {
        #     "chosen_ids":  [BOS, user..., chosen assistant..., EOS],
        #     "chosen_mask": [1, 1, ..., 1],
        #     "rejected_ids":  [BOS, user..., rejected assistant..., EOS],
        #     "rejected_mask": [1, 1, ..., 1],
        # }

    if dropped_identical:
        print(
            f"Dropped {dropped_identical} pairs whose chosen/rejected token ids were "
            f"identical after truncation to max_length={max_length}"
        )

    return Dataset.from_list(records)


def pad_sequences(sequences: List[List[int]], pad_value: int) -> torch.Tensor:
    """Right-pad variable-length Python lists and stack them as [B, T]."""
    max_len = max(len(sequence) for sequence in sequences)
    return torch.tensor(
        [sequence + [pad_value] * (max_len - len(sequence)) for sequence in sequences],
        dtype=torch.long,
    )


def collate_fn(batch: List[Dict], tokenizer: AutoTokenizer) -> Dict[str, torch.Tensor]:
    """Collate function that pads chosen and rejected sequences."""
    result = {
        "chosen_ids": pad_sequences(
            [x["chosen_ids"] for x in batch], tokenizer.pad_token_id
        ),
        "chosen_mask": pad_sequences([x["chosen_mask"] for x in batch], 0),
        "rejected_ids": pad_sequences(
            [x["rejected_ids"] for x in batch], tokenizer.pad_token_id
        ),
        "rejected_mask": pad_sequences([x["rejected_mask"] for x in batch], 0),
    }

    # Example for two preference records:
    # chosen lengths   = [120, 96]  -> chosen_ids.shape   = [2, 120]
    # rejected lengths = [104, 131] -> rejected_ids.shape = [2, 131]
    #
    # Chosen and rejected tensors are padded independently. The shorter second
    # chosen row gets 24 pad IDs and 24 mask zeros; the shorter first rejected
    # row gets 27 pad IDs and 27 mask zeros. Both still represent two pairs,
    # aligned by batch row: chosen[0] competes with rejected[0].
    return result


# Load tokenizer
config.dataset_name = "argilla/ultrafeedback-binarized-preferences-cleaned"
tokenizer = load_tokenizer(config.model_id)
if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token

# Build dataset
print(f"Building preference dataset with {config.samples} pairs...")
data = build_preference_dataset(tokenizer, config)
print(f"Dataset size: {len(data)} pairs")

if len(data) < (2 if config.val_ratio > 0 else 1):
    raise ValueError(
        "Too few usable preference pairs after tokenization; increase samples/max_length."
    )
if config.val_ratio > 0.0:
    splits = data.train_test_split(
        test_size=config.val_ratio,
        seed=config.seed,
        shuffle=True,
    )
    train_data = splits["train"]
    val_data = splits["test"]
else:
    train_data = data
    val_data = None

print(f"Train size: {len(train_data)} pairs")
if val_data is not None:
    print(f"Validation size: {len(val_data)} pairs")

train_loader = DataLoader(
    train_data,
    batch_size=config.batch_size,
    shuffle=True,
    drop_last=len(train_data) > config.batch_size,
    collate_fn=lambda b: collate_fn(b, tokenizer),
)

# train_loader now yields the four tensors consumed by RM training:
# chosen_ids [B, T_chosen], chosen_mask [B, T_chosen],
# rejected_ids [B, T_rejected], rejected_mask [B, T_rejected].
```

### Generation vs. Training Direction

- **RL Generation (Inference/Rollout):** **Left-pad** prompts so the active sequence ends align at index `[-1]`, ensuring batched autoregressive sampling attends to actual tokens rather than padding.
- **SFT & RL Loss Computation (Forward/Backward):** **Right-pad** (or pack sequences) so meaningful text begins at position `0`, preserving canonical position embeddings and standard lower-triangular causal masks.

## Minimum LLM training flow

### `grad_norm`

#### Loss, gradient, and gradient norm

- **Loss:** How poorly did the model predict the expected answer?
- **Gradient:** How would changing each model weight affect that loss? Each component is a partial derivative of the loss with respect to one weight.
- **Gradient norm:** One number describing the combined magnitude of all those gradients. The usual L2 norm treats every parameter's gradient as part of one long vector—not as separate norms to clip independently.

For model weights \(\theta_1,\ldots,\theta_n\):

$$
g=\nabla_\theta L=\left(\frac{\partial L}{\partial\theta_1},\ldots,\frac{\partial L}{\partial\theta_n}\right),\qquad
\|g\|_2=\sqrt{\sum_{i=1}^{n}g_i^2}.
$$

Loss measures prediction error; gradient norm measures local sensitivity to weight changes. They are different signals: a large loss does not necessarily imply a large gradient norm.

#### Gradient clipping is a cap, not normalization

**Gradient clipping caps unusually large gradient norms; it does not normalize every gradient to the same norm.** It rescales the entire gradient vector only when its norm exceeds the threshold. This limits unusually large gradients that could destabilize training while preserving their direction.

Let `max_grad_norm = 1.0`: a positive threshold, **not `-1.0`**. Suppose:

$$
g=(3,4),\qquad \|g\|_2=\sqrt{3^2+4^2}=5.
$$

Since \(5>1\), the shared scale factor is:

$$
\text{scale}=\frac{\text{max\_grad\_norm}}{\|g\|_2}=\frac{1}{5}=0.2.
$$

Every gradient component gets the same multiplier:

$$
g'=(3,4)\times0.2=(0.6,0.8),\qquad \|g'\|_2=1.
$$

But if \(g=(0.3,0.4)\), its norm is only \(0.5<1.0\). **Nothing changes:** \(g'=g\). Clipping never scales this smaller gradient up to norm 1. That is the difference from normalization.

#### From the gradient to a weight update

Imagine a model with only two parameters: \(a=2.5\) and \(b=4\). Its loss is a 3D surface: the x-axis is \(a\), the y-axis is \(b\), and the height is \(L(a,b)\).

At the current point, suppose:

$$
\nabla L=\left(\frac{\partial L}{\partial a},\frac{\partial L}{\partial b}\right)=(3,4).
$$

The gradient describes the direction and steepness of the loss surface. Locally, it points most steeply uphill, toward increasing \(a\) and \(b\); the negative gradient \((-3,-4)\) points most steeply downhill. Think of standing on a hill: to reduce the loss, take a sufficiently small step downhill.

Plain gradient descent subtracts the gradient, scaled by learning rate \(\eta\):

$$
a_{\mathrm{new}}=a-\eta\frac{\partial L}{\partial a},\qquad
b_{\mathrm{new}}=b-\eta\frac{\partial L}{\partial b}.
$$

With \(\eta=0.1\), **without clipping**:

$$
a_{\mathrm{new}}=2.5-0.1(3)=2.2,\qquad
b_{\mathrm{new}}=4-0.1(4)=3.6.
$$

**After clipping** the gradient to \((0.6,0.8)\), the same learning rate gives:

$$
a_{\mathrm{new}}=2.5-0.1(0.6)=2.44,\qquad
b_{\mathrm{new}}=4-0.1(0.8)=3.92.
$$

The gradient direction stays the same; the step becomes smaller. The threshold controls the gradient magnitude, while the learning rate controls how strongly plain gradient descent applies it. The skeleton below uses **AdamW**, which also uses moment estimates and weight decay, so its actual weight update is not simply \(-\eta g\).

#### Code intuition and the value you log

Conceptually, global norm clipping looks like this:

```python
# Pseudocode: one norm and one shared scale across all parameter gradients.
total_norm = norm(all_parameter_gradients)

if total_norm > max_grad_norm:
    scale = max_grad_norm / total_norm
    for p in model.parameters():
        if p.grad is not None:
            p.grad *= scale
```

In PyTorch, use it **after backward / gradient accumulation and before the optimizer update**:

```python
# The accumulation group's backward() calls have already populated p.grad.
grad_norm = torch.nn.utils.clip_grad_norm_(
    model.parameters(),
    max_norm=1.0,
)
optimizer.step()                      # Consume the clipped gradients.
optimizer.zero_grad(set_to_none=True) # Start the next group independently.
```

- **What changes:** `clip_grad_norm_` modifies `.grad` in place, across the model as one combined vector. It does not directly change the weights; `optimizer.step()` does that.
- **What is returned:** `grad_norm` is the total norm **before clipping**. In the `(3, 4)` example, the returned value is `5.0` even though the clipped norm is approximately `1.0`. In the small-gradient example it is `0.5`, and the gradients stay unchanged. See the [PyTorch implementation](https://github.com/pytorch/pytorch/blob/main/torch/nn/utils/clip_grad.py).
- **How to read it:** Compare the logged norm with `max_grad_norm` to see when clipping activates. Unlike `loss_ema`, which only smooths a reported metric, clipping changes the gradients used for training. In the skeleton, this norm is measured once per optimizer update, after the scaled micro-batch gradients have accumulated.

### Implementation

```python
def train(cfg):
    # Reproducibility and device: seed every RNG, then pick the requested accelerator.
    random.seed(cfg.seed)
    np.random.seed(cfg.seed)
    torch.manual_seed(cfg.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(cfg.seed)
    device = torch.device(
        cfg.device
        or (f"cuda:{cfg.model_device_id}" if torch.cuda.is_available()
            else "mps" if torch.backends.mps.is_available()
            else "cpu")
    )

    # Model and tokenizer: padding is needed because examples have different lengths.
    tokenizer = AutoTokenizer.from_pretrained(cfg.model_name)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(cfg.model_name).to(device)

    # Data loading was defined above; the training loop only consumes batches.
    dataloader = create_sft_dataloader(cfg, tokenizer)

    # Optimizer and schedule: both advance once per complete accumulation group.
    accum = cfg.gradient_accumulation_steps
    steps_per_epoch = len(dataloader) // accum
    total_steps = steps_per_epoch * cfg.num_epochs
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay
    )
    scheduler = get_scheduler(
        "cosine",
        optimizer=optimizer,
        num_warmup_steps=int(total_steps * cfg.warmup_ratio),
        num_training_steps=total_steps,
    )

    # Training state: gradients live on parameters; these scalars are for logging.
    model.train()
    optimizer.zero_grad(set_to_none=True)
    global_step = 0
    accumulated_loss = 0.0
    micro_batches = 0
    loss_ema = None

    # Epoch → mini-batch → loss → backward → optimizer step → loss aggregation.
    for epoch in range(cfg.num_epochs):
        for batch_idx, batch in enumerate(dataloader):
            batch = {name: tensor.to(device) for name, tensor in batch.items()}

            logits = model(
                input_ids=batch["input_ids"],
                attention_mask=batch["attention_mask"],
                use_cache=False,
            ).logits
            loss = F.cross_entropy(
                logits[:, :-1, :].contiguous().view(-1, logits.size(-1)),
                batch["labels"][:, 1:].contiguous().view(-1),
                ignore_index=-100,
            )

            if loss.isfinite():
                (loss / accum).backward()
                accumulated_loss += loss.item()
                micro_batches += 1

            if (batch_idx + 1) % accum == 0:
                # Clip accumulated gradients; record their norm before clipping.
                grad_norm = torch.nn.utils.clip_grad_norm_(
                    model.parameters(), max_norm=cfg.max_grad_norm
                )
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad(set_to_none=True)
                global_step += 1

                step_loss = accumulated_loss / max(1, micro_batches)
                # Mini-batches differ in examples, lengths, and difficulty, so
                # raw loss jumps from batch to batch. Smooth it for logging:
                # L_EMA,t = 0.9 * L_EMA,t-1 + 0.1 * L_t.
                # This metric does not change gradients or model weights.
                loss_ema = (
                    step_loss if loss_ema is None
                    else 0.9 * loss_ema + 0.1 * step_loss
                )
                accumulated_loss = 0.0
                micro_batches = 0

        # Drop an incomplete accumulation group so every update has the same size.
        optimizer.zero_grad(set_to_none=True)
        accumulated_loss = 0.0
        micro_batches = 0

    return model, {"global_step": global_step, "loss_ema": loss_ema}
```

## RL Training and Rollout

### Overview visual

```mermaid
sequenceDiagram
    participant T as train.py outer loop
    participant R as RolloutEngine
    participant P as Policy model
    participant V as PPO value model
    participant B as ReplayBuffer on CPU
    participant O as AdamW

    T->>R: iter(rollout_engine), then next()
    loop Until buffer reaches expected size
        R->>P: Generate completion(s), gradients disabled
        P-->>R: Prompt + sampled token IDs
        R->>P: Compute log_probs_old
        R->>V: Compute values_old
        R->>R: Score answers and compute advantages
        R->>B: Add Experience row(s), move tensors to CPU
    end

    R-->>T: yield replay_buffer
    Note over R,T: Rollout generator pauses here

    T->>T: Log average rewards to W&B

    loop Reuse same ReplayBuffer for ppo_epochs
        T->>B: Rebuild DataLoader and reshuffle rows
        loop Shuffled training micro-batches
            B-->>T: Move only this Experience to GPU/MPS
            T->>P: Recompute current log_probs under pi_theta
            T->>V: Recompute current values
            T->>T: Compare with fixed log_probs_old and clip loss
            T->>T: Backward with gradient accumulation
            T->>O: Step optimizer after accumulation
            O-->>P: pi_theta changes while pi_old stays fixed
        end
    end

    T->>R: next() resumes generator
    Note over R,P: Next rollouts use updated policy
```

### Multi-Epoch 原地复用 Buffer

一次 rollout 通常比一次训练 forward/backward 更昂贵，所以 PPO/GRPO 不会只读一次 `ReplayBuffer` 就立刻丢掉。训练主循环有两层：**外层产生一批新的 on-policy rollout 并冻结这一批数据的 $\pi_{\text{old}}$；内层在同一个 Buffer 上训练 `ppo_epochs` 次，每个 epoch 重新 shuffle。**

```python
# -----------------------------------------------------------------------------
# Outer loop: expensive rollout sampling creates new data and a new pi_old
# -----------------------------------------------------------------------------
for replay_buffer in rollout_engine:
    # The complete ReplayBuffer remains on CPU throughout all inner epochs.
    # Each Experience stores fixed rollout-time quantities such as:
    #   sequence_ids, action_mask, advantages, log_probs_old, values_old

    # -------------------------------------------------------------------------
    # Inner loop: reuse this exact Buffer in place for several update epochs
    # -------------------------------------------------------------------------
    for epoch in range(cfg.ppo_epochs):
        # Rebuild the sampler so the same rows receive a fresh random order.
        experience_sampler = DataLoader(
            replay_buffer.buffer,
            batch_size=cfg.train_batch_size,
            shuffle=True,
            collate_fn=join_experiences_batch,
        )

        for exp in experience_sampler:
            # Only the current micro-batch enters GPU/MPS memory; the full Buffer
            # stays on CPU. Its activations are released after backward/step.
            exp = exp.to(model.device)

            # pi_theta changes after every optimizer step, so these current-policy
            # log-probabilities must be recomputed on every visit to the row.
            log_probs = compute_log_probs(
                model, exp.sequence_ids, exp.attention_mask, cfg
            )

            # exp.log_probs_old never changes during these inner epochs.
            # ratio = exp(log pi_theta - log pi_old)
            # Before the first update it is ideally 1; later it drifts away from 1.
            loss = objective(log_probs=log_probs, experience=exp)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)

    # Only now is this Buffer retired. The outer loop resumes rollout generation
    # with the newly updated policy and records a fresh pi_old for the next Buffer.
```

#### 为什么带 Clip 的 PPO/GRPO 可以复用多次？

- **数据越来越“旧”：**第一次 inner update 前，$\pi_\theta$ 与生成 Buffer 的 $\pi_{\text{old}}$ 相同，故 $\rho=\exp(\log\pi_\theta-\log\pi_{\text{old}})\approx1$。每次 `optimizer.step()` 都会改变 $\pi_\theta$，但 Buffer 中的 `log_probs_old` 锁死不变，因此后续 epochs 的 $\rho$ 会逐渐偏离 $1$。
- **Importance sampling 先校正：**$\rho$ 衡量当前策略与采样策略对同一 token 的概率差异，让旧 Buffer 仍能为当前策略提供可解释的更新权重。
- **Clipping 再踩刹车：**当更新已经沿 advantage 希望的方向越过 $[1-\epsilon,1+\epsilon]$，PPO/GRPO 的 clipped objective 在该方向停止继续提供梯度，避免反复消费同一批样本时无限放大策略漂移。
- **Plain PG 不默认这样复用：**VPG、REINFORCE、RLOO 若没有 ratio clipping、KL trust region 或 early stopping，反复优化同一批 stale samples 没有自动刹车。因而在本页的简化实现中，它们通常只消费一次 Buffer；`ppo_epochs>1` 留给带显式 policy-change constraint 的目标。

显存生命周期可以记成：**整批 Buffer 跨 epochs 常驻 CPU；GPU/MPS 只暂存当前 micro-batch、当前模型参数和这次 forward/backward 的 activations。**多轮复用增加的是计算次数，而不是把同一 Buffer 复制 `ppo_epochs` 份。

### Rollout Engine

```python
class RolloutEngine:
    """
    可迭代的数据采样引擎：用于连接底层 DataLoader 与外部的训练循环。
    每次外层调用迭代（for buffer in engine），就会产出（yield）一个装满有效数据的 ReplayBuffer。
    """

    def __init__(self, cfg, dataset, model, tokenizer, ref_model=None, val_model=None):
        self.cfg = cfg
        self.dataset = dataset
        self.model = model
        self.tokenizer = tokenizer
        self.ref_model = ref_model
        self.val_model = val_model

        # 1. 设置自回归生成配置（RL Rollout 生成阶段强制要求左侧填充 Left-Padding）
        self.tokenizer.padding_side = "left"
        self.generation_config = GenerationConfig(
            temperature=cfg.temperature,
            top_p=cfg.top_p,
            do_sample=True,
            max_new_tokens=cfg.max_new_tokens,
            pad_token_id=self.tokenizer.pad_token_id,
        )

        # 2. 构造底层的 Prompt 批次加载器
        self._dataloader = DataLoader(
            dataset=dataset,
            batch_size=cfg.prompts_per_step,
            shuffle=True,
            collate_fn=lambda x: x,
        )

    def _iter_entries(self):
        """
        [水龙头生成器]：把 DataLoader 吐出来的二维 Batch 拍平成单个 entry。
        外层每 call 一次 next()，这里就 yield 出一个题目样本 entry。
        """
        for batch in self._dataloader:
            for entry in batch:
                yield entry

    def __iter__(self):
        """
        [流程总调度]：当外部执行 `for buffer in engine:` 时，首先触发该函数。
        """
        # Step 1: 实例化底层生成器，将其挂在 self._entries 上供后续反复读取
        self._entries = self._iter_entries()

        # Step 2: 驱动循环不断生产 Buffer
        while True:
            # 委派 _step() 去收集、推理并装满一个 ReplayBuffer
            buffer = self._step()

            # 如果数据耗尽，返回空 buffer，整场训练迭代结束
            if len(buffer) == 0:
                return

            # Step 3: [关键交货点] 挂起当前协程，将装满的 buffer 吐给外部训练循环做 backward
            # 外部做完单步训练后，代码会从这里唤醒，继续下一次 while True
            yield buffer

    def _step(self):
        """
        [蓄水池装填器]：负责不断索要 entry、推理采样、过滤，直到塞满一个 Step 规定的样本量。
        """
        # Rollout 采样阶段：模型全部进入 eval 模式，关闭训练图
        self.model.eval()
        if self.val_model is not None:
            self.val_model.eval()

        buffer = ReplayBuffer()
        # 计算一个 Buffer 所需收集的目标样本总数（例如: 1道题 * 4次采样 = 4条）
        expected = self.cfg.prompts_per_step * self.cfg.num_rollouts

        # 循环采水：只要合格的数据没凑齐，就一直循环
        while len(buffer) < expected:
            # 拧一下水龙头：向生成器索要下一个具体的题目 entry
            entry = next(self._entries, None)
            if entry is None:
                break  # 数据集读完了

            with torch.no_grad():
                # 拿当前题目做并发推理，得到候选回答并算好 Reward / Advantage
                exp = self._generate_experience(entry)
                # 质量过滤（如全对或全错样本无法产生组内对比信号，直接丢弃）
                exp = self._filter(exp)

            # 只有通过过滤的有效样本才被塞入 Buffer
            if exp is not None:
                buffer.add(exp)

        return buffer

    def _generate_experience(self, entry):
        """
        [组内采样与核心 RL 变量计算]：
        针对单道题目复制 G 份（Group Size），自回归生成多条回复，并计算 GRPO/PPO 所需的关键张量。
        """
        # 1. 将同一个 prompt 复制 num_rollouts 份，用于后续组内对抗（Group-level Comparison）
        entries = [entry for _ in range(self.cfg.num_rollouts)]
        prompts = [
            self.tokenizer.apply_chat_template(
                [{"role": "user", "content": e["question"]}],
                tokenize=False,
                add_generation_prompt=True,
            )
            for e in entries
        ]
        inputs = self.tokenizer(prompts, return_tensors="pt", padding=True).to(self.model.device)

        # 2. 策略模型自回归抽卡生成 (Rollout)
        sequence_ids = self.model.generate(**inputs, generation_config=self.generation_config)
        prompt_len = inputs["input_ids"].shape[1]
        completion_ids = sequence_ids[:, prompt_len:]
        completions = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True)

        # 3. 构造 action_mask：只把模型新生成的有效 token 标为 True（Prompt 和 Pad 全部排除）
        action_mask = torch.zeros_like(sequence_ids, dtype=torch.bool)
        action_mask[:, prompt_len:] = True
        action_mask[sequence_ids == self.tokenizer.pad_token_id] = False
        action_mask = action_mask[:, 1:]
        attention_mask = sequence_ids != self.tokenizer.pad_token_id

        # 4. 评估答案，获得整句标量奖励 Rewards
        rewards = compute_rewards(entries, completions, self.dataset, self.cfg)

        # 5. 计算并缓存基准对数概率（pi_old，用于后续训练算重要性采样 ratio）
        log_probs_old = compute_log_probs(self.model, sequence_ids, attention_mask, self.cfg)
        log_probs_ref = compute_log_probs(self.ref_model, sequence_ids, attention_mask, self.cfg)
        values_old = compute_values(self.val_model, sequence_ids, attention_mask, self.cfg)

        # 6. 计算 Advantages（在 GRPO 里是组内打分标准化；在 PPO 里是 GAE）
        rewards = apply_reward_kl(rewards, log_probs_old, log_probs_ref, action_mask, self.cfg)
        advantages = compute_advantages(rewards, action_mask, values_old, self.cfg)

        # 7. 打包成 Experience 对象，并转移到 CPU 节省 GPU/MPS 显存
        return Experience(
            sequence_ids=sequence_ids,
            attention_mask=attention_mask,
            action_mask=action_mask,
            rewards=rewards,
            log_probs_old=log_probs_old,
            log_probs_ref=log_probs_ref,
            values_old=values_old,
            advantages=advantages,
        ).to("cpu")

    def _filter(self, exp):
        """样本筛选器：若组内所有回答得分完全相同（全满分或全零分），组内方差为0，失去对比学习价值，直接丢弃"""
        if self.cfg.loss == "dapo":
            correctness = exp.rewards["correctness"]
            all_min = (correctness == self.cfg.accuracy_min_reward).all()
            all_max = (correctness == self.cfg.accuracy_max_reward).all()
            if all_min or all_max:
                return None
        return exp
```

### Policy Optimization Loop

```python
# ==============================================================================
# 强化学习主训练循环 (Entry Point): 串联 Rollout 与 Policy Optimization
# ==============================================================================

# 1. 实例化采样引擎（它是一个 Python 生成器 / Iterator）
rollout_engine = RolloutEngine(
    cfg=cfg,
    dataset=dataset,
    model=model,
    tokenizer=tokenizer,
    ref_model=ref_model,
    val_model=val_model,
)

# 2. 外层循环：昂贵的 Rollout 采样。
# 每次向 RolloutEngine 索要一个新 ReplayBuffer，并记录这一轮固定的 pi_old。
for replay_buffer in rollout_engine:
    # --------------------------------------------------------------------------
    # Phase 1: 准备训练环境
    # --------------------------------------------------------------------------
    model.train()
    if val_model:
        val_model.train()

    # --------------------------------------------------------------------------
    # Phase 2: 内层循环原地复用同一个 Buffer cfg.ppo_epochs 次。
    # Buffer 整体常驻 CPU；log_probs_old / values_old / advantages 始终锁死为
    # rollout-time 数值，但当前模型 pi_theta 会在每次 optimizer.step() 后变化。
    # --------------------------------------------------------------------------
    for epoch in range(cfg.ppo_epochs):
        # 每个 epoch 重建 DataLoader，让同一批 Experience 获得新的随机顺序。
        experience_sampler = DataLoader(
            dataset=replay_buffer.buffer,
            batch_size=cfg.train_batch_size,
            shuffle=True,
            drop_last=False,
            collate_fn=join_experiences_batch,
        )
        optimizer.zero_grad(set_to_none=True)

        for idx, exp in enumerate(experience_sampler):
            # 1. 只把当前 micro-batch 移入 GPU/MPS；完整 Buffer 仍留在 CPU。
            exp = exp.to(model.device)

            # 2. 每次访问都用最新 pi_theta 重新计算 log_probs / values。
            # 第一个 optimizer step 之前，若权重和前向设置相同，ratio 理论上为 1；
            # 之后 pi_theta 持续变化，而 exp.log_probs_old 在所有 epochs 中保持不变。
            log_probs = compute_log_probs(model, exp.sequence_ids, exp.attention_mask, cfg)
            values = compute_values(val_model, exp.sequence_ids, exp.attention_mask, cfg)

            # 3. PPO / GRPO: ratio = exp(log_probs - exp.log_probs_old)，
            # 再由 clipped objective 限制同一 Buffer 被反复使用时的策略漂移。
            loss = objective(log_probs=log_probs, experience=exp, values=values)

            # 4. 梯度累积归一化：尾部 batch 可能不足，动态除数保持平均正确。
            divisor = accumulation_divisor(idx, len(experience_sampler), cfg.batch_acc)
            scaled_loss = loss / divisor
            scaled_loss.backward()

            # 5. 达到累积步数或当前 epoch 的最后一个 batch 时更新参数。
            is_accum_step = (idx + 1) % cfg.batch_acc == 0
            is_last_batch = (idx + 1) == len(experience_sampler)
            if is_accum_step or is_last_batch:
                clip_grad_norm_(params, max_norm=cfg.max_norm)
                optimizer.step()  # pi_theta 改变；Buffer 中的 pi_old 不变
                optimizer.zero_grad(set_to_none=True)

# 同一个 ReplayBuffer 完成全部 ppo_epochs 后才退出内层循环并释放。
# 外层再次唤醒 rollout_engine，用更新后的策略生成新数据和新的 pi_old。
```


## PyTorch Coding

### `*` 与 `@` 的区别和应用

- **`@`（matrix multiplication）：**沿共享的内维度做乘加并改变张量表示，常用于让不同 features、tokens 或权重发生线性组合。
- **`*`（element-wise multiplication）：**对应位置逐元素相乘，通常保持原 shape；两个张量 shape 不完全相同时，必须满足 PyTorch broadcasting 规则。

| 运算 | 典型形状流动 | 常见 LLM 步骤 / 场景 |
|---|---|---|
| **`@` (Matmul)** | `[B, T, D] @ [D, D'] -> [B, T, D']` | Attention 交互（$QK^T$、$AV$）、MLP 线性投影、LM Head 预测 logits |
| **`*` (Element-wise)** | `[B, T] * [B, T] -> [B, T]` | RLHF / GRPO policy loss 赋权（$-A\times\log p$）、padding mask 过滤、SwiGLU 门控、RoPE 旋转 |

最简判断方法：`@` 会对某个维度**求和并混合信息**；`*` 只把对应位置的值**缩放或门控**。

### 维度变换：`reshape` 与 `None`

#### 把 `[B, G]` 展平并增加最后一维

```python
advantages = advantages.reshape(-1, 1)
```

- **`-1`（自动推断展平后的长度）：**原来的 `[B, G]` 被拍平成一个长度为 $B\times G$ 的维度。若 $B=2$、$G=4$，长度就是 $8$。
- **`1`（增加 singleton dimension）：**最后一维被固定为 $1$。
- **形状变化：**

```math
\text{advantages}: [B,G]
\xrightarrow{\texttt{reshape(-1, 1)}}
[B\times G,1].
```

本例中，`advantages` 从 `[2, 4]` 变为 **`[8, 1]`**。`-1` 不是“删除一个维度”，而是让 PyTorch 根据元素总数自动算出这一维应为 $8$；元素的数量和顺序都没有改变。

#### 用 `None` 插入一个大小为 1 的维度

```python
sequence_ratio = sequence_ratio[:, None]
```

- `:`：保留第 0 维的所有 $B\times G$ 个元素。
- `None`：在第 1 维插入一个大小为 $1$ 的新维度；它等价于 `sequence_ratio.unsqueeze(1)`。
- **形状变化：**`[B*G] -> [B*G, 1]`，例如 `[8] -> [8, 1]`。
- **目的：**若后面的 token-level tensor 形状为 `[B*G, T]`，那么 `[B*G, 1]` 的 sequence-level ratio 或 advantage 可以沿 token 维自动广播，把同一个 sequence-level 值应用到该序列的全部 $T$ 个 token。

`reshape(-1, 1)` 与 `[:, None]` 最终都可能得到 `[B*G, 1]`，但起点不同：前者把 `[B,G]` 的多个维度合并；后者从 `[B*G]` 开始，只插入新维度。

### Broadcasting 机制

假设每个 prompt 采样 $G$ 个 responses，batch 内共有 $B$ 个 prompts：

- `rewards`：`[B, G]`，本例为 `[2, 4]`。
- `group_mean`：`[B, 1]`，本例为 `[2, 1]`；每个 prompt 对应一个组内平均分。
- `group_std`：`[B, 1]`，本例为 `[2, 1]`；每个 prompt 对应一个组内标准差。

#### 组内归一化计算 Advantage

```python
advantages = (rewards - group_mean) / (group_std + 1e-4)
```

Broadcasting 从每一维的末尾向前比较 shape：两个维度相等，或者其中一个维度为 `1`，就可以兼容。因此：

```text
rewards:    [2, 4]
group_mean: [2, 1]  -> logically expands to [2, 4]
group_std:  [2, 1]  -> logically expands to [2, 4]
result:     [2, 4]
```

这里 `[2,1]` 的第二维会沿 $G=4$ 个 responses 逻辑扩展：同一个 prompt 的 mean 和 std 分别作用于该组的四个 rewards，完成组内减均值与除以标准差。PyTorch 通常通过 stride/view 语义实现这种广播，并不需要真的复制出四份数据。

- **输出形状：**`advantages` 仍为 `[B, G]`，本例是 **`[2, 4]`**。
- **`1e-4` 的作用：**避免某组所有 rewards 相同、`group_std = 0` 时发生除零。
- **不会广播的例子：**`[2,4]` 与 `[2]` 不能按“每个 prompt 一个值”的意图直接运算，因为末维 `4` 和 `2` 不兼容；应先把 `[2]` 变为 `[2,1]`。