14. 规模扩展与效率

#@save
import torch
from torch.utils.checkpoint import checkpoint
import agentrl as arl

import copy
torch.set_num_threads(4)                       # the shared check host is heavily loaded
tok = arl.build_tokenizer(arl.make_tasks('all'))
train_tasks, heldout_tasks = arl.make_tasks('train'), arl.make_tasks('heldout')
rl_base = arl.make_rl_base(tok)
cfg = arl.RLConfig()
trajs = arl.sample_groups(rl_base, lambda: arl.make_env(cfg), train_tasks[:8], tok, cfg,
                          torch.Generator().manual_seed(0))                     # 64 real episodes

lengths = torch.tensor([len(t.ids) for t in trajs])                            # [64] tokens per episode
real, padded = lengths.sum().item(), len(lengths) * lengths.max().item()
q = lengths.float().quantile(torch.tensor([0.5, 0.9]))
print(f'episode tokens: min {lengths.min()} | median {q[0]:.0f} | p90 {q[1]:.0f} | max {lengths.max()}')
print(f'padded batch: {padded} positions for {real} real tokens -> {1 - real / padded:.0%} of compute is padding')

episode tokens: min 10 | median 18 | p90 79 | max 149
padded batch: 9536 positions for 2472 real tokens -> 74% of compute is padding

#@save
def pack_sequences(trajs, max_tokens, pad_id=0):
    """First-fit-decreasing packing of Trajectory objects into rows of max_tokens -> dict of [R, L] / [R, L-1]."""
    rows, used = [], []
    for i in sorted(range(len(trajs)), key=lambda i: -len(trajs[i].ids)):
        n = len(trajs[i].ids)
        if n > max_tokens:
            raise ValueError(f'trajectory of {n} tokens does not fit in a row of {max_tokens}')
        r = next((r for r, u in enumerate(used) if u + n <= max_tokens), None)
        if r is None:
            rows.append([]), used.append(0)
            r = len(rows) - 1
        rows[r].append(i)
        used[r] += n
    R = len(rows)
    out = dict(ids=torch.full((R, max_tokens), pad_id), position_ids=torch.zeros(R, max_tokens, dtype=torch.long),
               segment_ids=torch.full((R, max_tokens), -1), action_mask=torch.zeros(R, max_tokens - 1),
               sample_logp=torch.zeros(R, max_tokens - 1), rows=rows)
    for r, members in enumerate(rows):
        offset = 0
        for k, i in enumerate(members):
            t, n = trajs[i], len(trajs[i].ids)
            out['ids'][r, offset:offset + n] = t.ids
            out['position_ids'][r, offset:offset + n] = torch.arange(n)       # positions restart per sequence
            out['segment_ids'][r, offset:offset + n] = k
            out['action_mask'][r, offset:offset + n - 1] = t.action_mask     # last position predicts nothing
            out['sample_logp'][r, offset:offset + n - 1] = t.sample_logp
            offset += n
    return out

#@save
def packed_logits(model, ids, position_ids, segment_ids):
    """TinyAgent forward on packed rows with a block-diagonal causal mask: ids [R, L] -> logits [R, L, V]."""
    L = ids.shape[1]
    h = model.tok_emb(ids) + model.pos_emb(position_ids)                             # [R, L, D]
    causal = torch.ones(L, L, dtype=torch.bool).tril()                               # [L, L]
    allowed = causal & (segment_ids[:, :, None] == segment_ids[:, None, :])          # [R, L, L] same segment
    allowed |= torch.eye(L, dtype=torch.bool)                                        # padding sees itself
    mask = torch.zeros(allowed.shape).masked_fill(~allowed, float('-inf'))
    mask = mask.repeat_interleave(model.blocks.layers[0].self_attn.num_heads, dim=0) # [R * heads, L, L]
    return model.head(model.norm(model.blocks(h, mask=mask)))

packed = pack_sequences(trajs, 256, tok.pad_id)
print(f'{len(trajs)} episodes -> {len(packed["rows"])} rows of 256 | fill {(packed["segment_ids"] >= 0).float().mean():.1%}')
mask = packed['action_mask']                                                         # [R, L-1]
with torch.no_grad():
    alone = arl.pad_rows([arl.sequence_log_probs(rl_base, t.ids[None])[0] for t in trajs], 255, 0.0)
    reference = torch.zeros_like(mask)                                               # alone, in packed layout
    for r, members in enumerate(packed['rows']):
        offset = 0
        for i in members:
            n = len(trajs[i].ids)
            reference[r, offset:offset + n - 1] = alone[i, :n - 1]
            offset += n
    variants = {'causal mask, running positions': (torch.arange(256).expand_as(packed['ids']), 0 * packed['segment_ids']),
                'causal mask, restarted positions': (packed['position_ids'], 0 * packed['segment_ids']),
                'block mask, restarted positions': (packed['position_ids'], packed['segment_ids'])}
    for name, (pos, seg) in variants.items():
        logp = arl.selected_log_probs(packed_logits(rl_base, packed['ids'], pos, seg)[:, :-1], packed['ids'][:, 1:])
        err = ((logp - reference).abs() * mask).max().item()
        print(f'{name:<34} max |packed - alone| on policy tokens: {err:.1e} nats')

64 episodes -> 10 rows of 256 | fill 96.6%
causal mask, running positions     max |packed - alone| on policy tokens: 6.4e+00 nats
causal mask, restarted positions   max |packed - alone| on policy tokens: 1.7e+00 nats
block mask, restarted positions    max |packed - alone| on policy tokens: 1.4e-06 nats

#@save
def token_budget_batches(lengths, max_tokens):
    """Group sequence indices into microbatches whose padded size (count * longest) <= max_tokens -> list of lists."""
    order = sorted(range(len(lengths)), key=lambda i: -int(lengths[i]))
    batches, current = [], []
    for i in order:                                   # longest first: the first member sets the padded width
        if current and (len(current) + 1) * int(lengths[current[0]]) > max_tokens:
            batches.append(current)
            current = []
        current.append(i)
    return batches + ([current] if current else [])

gen = torch.Generator().manual_seed(0)
agent_lengths = (8192 * torch.exp(0.8 * torch.randn(512, generator=gen))).clamp(256, 65_536).long()   # [512]
def describe(name, batches):
    sizes = [len(b) * int(agent_lengths[b].max()) for b in batches]                 # padded positions
    waste = 1 - agent_lengths.sum().item() / sum(sizes)
    print(f'{name:<30} | {len(batches):4d} microbatches | padding {waste:5.1%} | largest {max(sizes) / 1e3:6.0f}k tokens')

random_order = torch.randperm(512, generator=gen)
describe('8 per microbatch, random', [random_order[i:i + 8] for i in range(0, 512, 8)])
describe('8 per microbatch, sorted', [agent_lengths.argsort()[i:i + 8] for i in range(0, 512, 8)])
describe('token budget 128k, sorted', [torch.tensor(b) for b in token_budget_batches(agent_lengths, 131_072)])

8 per microbatch, random       |   64 microbatches | padding 59.3% | largest    524k tokens
8 per microbatch, sorted       |   64 microbatches | padding  3.6% | largest    524k tokens
token budget 128k, sorted      |   51 microbatches | padding  3.0% | largest    131k tokens

V = 151_936
print(f'{"model":<11} {"D":>5} {"V / D":>6} | fp32 logits per token | logits of a 32k-token microbatch')
for name, D in [('Qwen3-0.6B', 1024), ('Qwen3-8B', 4096), ('Qwen3-32B', 5120)]:
    print(f'{name:<11} {D:5d} {V / D:6.0f} | {4 * V / 1e6:17.2f} MB | {4 * V * 32_768 / 1e9:13.1f} GB')

model           D  V / D | fp32 logits per token | logits of a 32k-token microbatch
Qwen3-0.6B   1024    148 |              0.61 MB |          19.9 GB
Qwen3-8B     4096     37 |              0.61 MB |          19.9 GB
Qwen3-32B    5120     30 |              0.61 MB |          19.9 GB

#@save
def checkpointed_log_probs(hidden, weight, ids, chunk_size=1024):
    """Chunked selected-token log-probs, recomputed in backward: hidden [B, T, D], weight [V, D], ids [B, T] -> [B, T]."""
    def chunk_logp(h, w, i):
        return arl.selected_log_probs(h @ w.T, i)                                 # [B, C]; logits live only here
    pieces = [checkpoint(chunk_logp, hidden[:, s:s + chunk_size], weight, ids[:, s:s + chunk_size],
                         use_reentrant=False) for s in range(0, hidden.shape[1], chunk_size)]
    return torch.cat(pieces, dim=1)

from transformers import AutoConfig, AutoModelForCausalLM
device = 'cuda'                                                   # peak-memory counters need the check host's GPU
weight = AutoModelForCausalLM.from_pretrained('Qwen/Qwen3-0.6B', dtype=torch.bfloat16).lm_head.weight
weight = weight.detach().to(device).requires_grad_(True)                          # [V, D] = [151936, 1024]

def peak_memory(logp_fn, T):
    """Peak GPU memory (GB) of logp_fn's forward + backward for T tokens, and the log-probs it returned."""
    gen = torch.Generator(device=device).manual_seed(0)
    hidden = torch.randn(1, T, weight.shape[1], device=device, dtype=torch.bfloat16, generator=gen).requires_grad_()
    ids = torch.randint(weight.shape[0], (1, T), device=device, generator=gen)
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    start = torch.cuda.memory_allocated()
    logp = logp_fn(hidden, weight, ids)                                            # [1, T]
    logp.sum().backward()
    torch.cuda.synchronize()
    weight.grad = None
    return (torch.cuda.max_memory_allocated() - start) / 1e9, logp.detach()

methods = {'full logits': lambda h, w, i: arl.selected_log_probs(h @ w.T, i),
           'chunked (Ch. 3)': lambda h, w, i: arl.chunked_log_probs(h, w, i, chunk_size=1024),
           'chunked + recompute': lambda h, w, i: checkpointed_log_probs(h, w, i, chunk_size=1024)}
for T in [4096, 16_384]:
    results = {name: peak_memory(fn, T) for name, fn in methods.items()}
    agree = max((r[1] - results['full logits'][1]).abs().max().item() for r in results.values())
    print(f'T = {T:6d} | ' + ' | '.join(f'{n} {r[0]:4.1f} GB' for n, r in results.items()) + f' | max diff {agree:.0e}')

T =   4096 | full logits  7.5 GB | chunked (Ch. 3)  3.7 GB | chunked + recompute  2.2 GB | max diff 0e+00
T =  16384 | full logits 29.9 GB | chunked (Ch. 3) 11.2 GB | chunked + recompute  2.2 GB | max diff 0e+00

batch = arl.collate(trajs, tok.pad_id)
advantages = arl.compute_advantages(batch, cfg)                                     # [64]
def accumulated_grad(model, chunks, normalize):
    """Sum of microbatch gradients of the REINFORCE loss; 'local' = token_mean per microbatch / K, 'global' = sum / N."""
    model.zero_grad()
    n_global = batch['action_mask'].sum()
    for rows in chunks:
        logp, m = arl.sequence_log_probs(model, batch['ids'][rows]), batch['action_mask'][rows]
        if normalize == 'local':
            loss = arl.reinforce_loss(logp, advantages[rows], m) / len(chunks)
        else:
            loss = arl.reinforce_loss(logp, advantages[rows], m, mode='constant', norm_const=n_global)
        loss.backward()
    return torch.cat([p.grad.flatten() for p in model.parameters()])

model = copy.deepcopy(rl_base)
full = accumulated_grad(model, [torch.arange(64)], 'local')
by_length = batch['action_mask'].sum(-1).argsort()
for split, chunks in [('random', torch.randperm(64, generator=torch.Generator().manual_seed(0)).chunk(4)),
                      ('length-sorted', by_length.chunk(4))]:
    tokens = [int(batch['action_mask'][c].sum()) for c in chunks]
    for normalize in ['local', 'global']:
        g = accumulated_grad(model, chunks, normalize)
        print(f'{split:<13} {normalize:<6} | policy tokens per microbatch {tokens} | relative error'
              f' {(g - full).norm() / full.norm():.1e} | cosine {torch.cosine_similarity(g, full, 0):.4f}')

random        local  | policy tokens per microbatch [159, 208, 185, 191] | relative error 8.2e-02 | cosine 0.9967
random        global | policy tokens per microbatch [159, 208, 185, 191] | relative error 1.3e-07 | cosine 1.0000
length-sorted local  | policy tokens per microbatch [96, 112, 180, 355] | relative error 5.2e-01 | cosine 0.8782
length-sorted global | policy tokens per microbatch [96, 112, 180, 355] | relative error 1.3e-07 | cosine 1.0000

#@save
def filter_by_pass_rate(tasks, pass_rates, low=0.0, high=1.0):
    """Tasks whose estimated pass rate lies strictly inside (low, high)."""
    return [task for task, p in zip(tasks, pass_rates) if low < p < high]

#@save
def train_curriculum(model, tasks, tok, cfg, steps, low=0.0, high=1.0, revisit=1):
    """train_rl sampling only tasks inside filter_by_pass_rate (plus `revisit` others) -> history of dicts."""
    torch.manual_seed(cfg.seed)
    gen = torch.Generator().manual_seed(cfg.seed)
    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)
    estimate, history, G = {t.name: 0.5 for t in tasks}, [], cfg.group_size
    for step in range(steps):
        pool = filter_by_pass_rate(tasks, [estimate[t.name] for t in tasks], low, high)
        excluded = [t for t in tasks if t not in pool]
        again = [excluded[i] for i in torch.randperm(len(excluded), generator=gen)[:revisit].tolist()]
        kept, sampled = arl.sample_step(model, lambda: arl.make_env(cfg), pool + again, tok, cfg, gen)
        flat = 0
        for g in range(0, len(sampled), G):
            rewards = [t.reward for t in sampled[g:g + G]]
            estimate[sampled[g].info['task']] = sum(rewards) / G               # last group's pass rate
            flat += len(set(rewards)) == 1
        batch = arl.collate(kept, tok.pad_id)
        stats = arl.rl_update(model, optimizer, batch, arl.compute_advantages(batch, cfg), cfg, None, gen)
        history.append(dict(step=step, pool=len(pool), n_episodes=len(sampled),
                            zero_adv_frac=flat / (len(sampled) // G), **arl.behavior_stats(sampled), **stats))
    return history

curriculum = {}
for arm, (low, high) in {'all tasks': (-1.0, 2.0), 'filter (0, 1)': (0.0, 1.0)}.items():
    for seed in [0, 1]:
        model = copy.deepcopy(rl_base)
        hist = curriculum[arm, seed] = train_curriculum(model, train_tasks, tok, arl.RLConfig(seed=seed), 60, low, high)
        held = arl.evaluate(model, heldout_tasks, tok, 16, generator=torch.Generator().manual_seed(seed), cfg=cfg)
        sampled = [sum(h['n_episodes'] for h in hist[:k]) for k in (20, 40, 60)]
        zero = ' '.join(f'{sum(h["zero_adv_frac"] for h in hist[i:i + 20]) / 20:.2f}' for i in range(0, 60, 20))
        print(f'{arm:<14} seed {seed} | episodes after 20/40/60 steps {sampled} | zero-adv share {zero}'
              f' | held-out {held["pass_rate"]:.3f}')

all tasks      seed 0 | episodes after 20/40/60 steps [1280, 2560, 3840] | zero-adv share 0.20 0.56 0.68 | held-out 0.756
all tasks      seed 1 | episodes after 20/40/60 steps [1280, 2560, 3840] | zero-adv share 0.14 0.39 0.61 | held-out 0.819
filter (0, 1)  seed 0 | episodes after 20/40/60 steps [1272, 1728, 2088] | zero-adv share 0.18 0.40 0.47 | held-out 0.731
filter (0, 1)  seed 1 | episodes after 20/40/60 steps [1280, 1808, 2184] | zero-adv share 0.17 0.43 0.38 | held-out 0.769

kv_free = kv_cached = 0
for t in trajs:
    policy_positions = t.action_mask.nonzero().flatten() + 1                    # index of each sampled token in ids
    kv_free += int(policy_positions.sum())                                      # a decode step reruns its whole prefix
    kv_cached += len(t.ids)                                                     # each position computed once
print(f'positions processed while sampling 64 episodes: no KV cache {kv_free:,} | KV cache {kv_cached:,}'
      f' | ratio {kv_free / kv_cached:.1f}x')

positions processed while sampling 64 episodes: no KV cache 22,661 | KV cache 2,472 | ratio 9.2x

from pathlib import Path
snapshot = next((Path.home() / '.cache/huggingface/hub/models--Qwen--Qwen3-8B/snapshots').iterdir())
qcfg = AutoConfig.from_pretrained(snapshot)                                     # local config.json only
N = 8.2e9                                                                        # parameters
kv_bytes = 2 * qcfg.num_hidden_layers * qcfg.num_key_value_heads * qcfg.head_dim * 2   # K and V, bf16, per token
BANDWIDTH, PEAK, MFU = 4.8e12, 989e12, 0.4                                      # H200; MFU is an assumption
print(f'Qwen3-8B KV cache: {kv_bytes / 1e3:.0f} KB per token of context')
def decode_tokens_per_s(b, c):
    return b / ((2 * N + b * c * kv_bytes) / BANDWIDTH)
for b, c in [(64, 2048), (64, 16_384), (256, 16_384)]:
    print(f'decode, batch {b:3d}, context {c:6d}: {decode_tokens_per_s(b, c):7.0f} tokens/s per GPU')
print(f'training (6N FLOPs per token, MFU {MFU}): {MFU * PEAK / (6 * N):7.0f} tokens/s per GPU')

Qwen3-8B KV cache: 147 KB per token of context
decode, batch  64, context   2048:    8598 tokens/s per GPU
decode, batch  64, context  16384:    1796 tokens/s per GPU
decode, batch 256, context  16384:    1935 tokens/s per GPU
training (6N FLOPs per token, MFU 0.4):    8041 tokens/s per GPU

episodes, total_tok, policy_tok, context = 512, 24_000, 6_000, 16_384     # per step; assumed episode shape
steps, busy = 300, 0.5                                                      # assumed run length; engine busy fraction
train_h = 6 * N * episodes * total_tok / (MFU * PEAK) / 3600                # forward + backward on every token
logp_h = 2 * N * episodes * total_tok / (MFU * PEAK) / 3600                 # one extra forward: old log-probs
prefill_h = 2 * N * episodes * (total_tok - policy_tok) / (MFU * PEAK) / 3600   # prompt and tool tokens, prefilled once
decode_h = episodes * policy_tok / decode_tokens_per_s(64, context) / 3600 / busy
parts = dict(train=train_h, old_logp=logp_h, prefill=prefill_h, decode=decode_h)
print(' | '.join(f'{k} {v:5.2f}' for k, v in parts.items()), f'| GPU-hours per step {sum(parts.values()):.2f}')
print(f'{steps} steps: {steps * sum(parts.values()):,.0f} GPU-hours; decode share {decode_h / sum(parts.values()):.0%}')

train  0.42 | old_logp  0.14 | prefill  0.11 | decode  0.95 | GPU-hours per step 1.62
300 steps: 487 GPU-hours; decode share 59%

tokens = torch.tensor([batch['action_mask'][c].sum().item() for c in by_length.chunk(4)])
distortion = tokens.sum() / (len(tokens) * tokens)                       # weight under local / weight under global
print('policy tokens per microbatch:', tokens.int().tolist(), '| per-token weight vs correct:',
      [round(x, 2) for x in distortion.tolist()])

policy tokens per microbatch: [96, 112, 180, 355] | per-token weight vs correct: [1.93, 1.66, 1.03, 0.52]

mask = packed['action_mask']
for name, (pos, seg) in [('leaky packing', variants['causal mask, restarted positions']),
                         ('correct packing', variants['block mask, restarted positions'])]:
    with torch.no_grad():
        logp = arl.selected_log_probs(packed_logits(rl_base, packed['ids'], pos, seg)[:, :-1], packed['ids'][:, 1:])
    ratio = ((logp - packed['sample_logp']).exp() - 1).abs() * mask
    print(f'{name:<16} share of policy tokens with |pi_trainer / mu - 1| > 0.2: {(ratio > 0.2).sum() / mask.sum():.1%}')

leaky packing    share of policy tokens with |pi_trainer / mu - 1| > 0.2: 5.8%
correct packing  share of policy tokens with |pi_trainer / mu - 1| > 0.2: 0.0%

for arm in ['all tasks', 'filter (0, 1)']:
    wasted = sum(h['zero_adv_frac'] * h['n_episodes'] for s in [0, 1] for h in curriculum[arm, s])
    total = sum(h['n_episodes'] for s in [0, 1] for h in curriculum[arm, s])
    print(f'{arm:<14} episodes in zero-advantage groups: {wasted:.0f} of {total} ({wasted / total:.0%})')

all tasks      episodes in zero-advantage groups: 3304 of 7680 (43%)
filter (0, 1)  episodes in zero-advantage groups: 1136 of 4272 (27%)

longest = int(agent_lengths.argmax())
fixed = next(b for b in [random_order[i:i + 8] for i in range(0, 512, 8)] if longest in b.tolist())
budget = next(b for b in token_budget_batches(agent_lengths, 131_072) if longest in b)
print(f'longest trajectory: {agent_lengths[longest]:,} tokens; median {int(agent_lengths.median()):,}')
print(f'fixed 8 per microbatch: its microbatch pads to {len(fixed) * int(agent_lengths[fixed].max()):,} tokens')
print(f'token budget: its microbatch holds {len(budget)} sequence(s), {len(budget) * int(agent_lengths[budget].max()):,} tokens')

longest trajectory: 65,536 tokens; median 8,414
fixed 8 per microbatch: its microbatch pads to 524,288 tokens
token budget: its microbatch holds 2 sequence(s), 131,072 tokens