vh
7b5fd91d3c
Run-01 was killed at step 19 by operator instruction to root-cause before spending a ~13.9h window. Two independent methods now agree on where the step time went, and neither was the hypothesis the consult panel converged on. Scaling fit (3 points, 2 params, residuals <3ms over an 8x range): A = 6.87e-4 s/token, B = 8.85e-8 s/token^2 quadratic share 20.9% @ w=2048 -> 67.8% @ w=16384 No fixed term was needed, which refutes launch-bound outright. Profiler kernel table (device rows only): attention 22,835.8 ms 65.2% fmha_cutlass*_sm80 dense GEMM 2,774.0 ms 7.9% other 5,739.0 ms 16.4% The attention kernels are sm80 — Ampere-generation CUTLASS running on an sm_120 Blackwell card, with the forward on the gmem fallback tier. That is the mechanism behind 100% SM utilisation at 27 of 304 available TFLOPS. Correctness cleared separately: the sliding mask asserts at max 1024 allowed/row, so the 25 windowed layers were genuinely windowed. The same probe found that right-padding is what pins the 5 global layers to an explicit 4D mask and off the is_causal fast path — measured at 9.4% slower for 24% less loss work at fixed width. The largest available win is not the attention kernel. The corpus is 29.9% padding, and bucket-to-pair + shuffle-to-mix takes it to 0.0% for >=35.5% wall clock, no new dependency, unchanged peak memory. Bucket size turned out not to be a diversity knob — roots per accumulation window are flat across a 256x range, so the global micro-batch shuffle does that work alone and the bucket should be tight. Adds docs/pfi/training-throughput-playbook.md as the durable model-agnostic home (sibling to the quantization playbook), the four probes under scripts/training-probes/ with raw output kept for re-derivation, and a §6 to the sizing doc carrying the Gemma-4-specific numbers and round-2 restart parameters. Measured negatives recorded so they are not re-chased: grouped_mm (0.9% slower, and MoE is only 7.9% of the step), CUDA graphs / torch.compile over the expert loop (no fixed cost to amortise), liger fused CE (~1-3% lever), FA4 on sm_120. Round-1 state preserved: 609MB encode cache, order manifest, truncation report, resume script. No checkpoints — it died at step 19 and the first was due at 100, so the lora_B inert-adapter gate never ran and moves to the restart.
76 lines
3.0 KiB
Python
76 lines
3.0 KiB
Python
"""Step 0 - assert the sliding mask band structure, and record which path
|
|
mask creation actually takes under the run config (attn_implementation=sdpa).
|
|
|
|
Correctness gate: transformers can SILENTLY skip mask creation and pass
|
|
attention_mask=None, which would make the 25 sliding layers do full causal
|
|
attention - a different model from the one vLLM serves. This converts
|
|
"probably fine because we are slow" into a measurement.
|
|
|
|
CPU only. No weights. No GPU.
|
|
"""
|
|
import torch
|
|
from transformers import AutoConfig
|
|
from transformers.masking_utils import (
|
|
create_causal_mask, create_sliding_window_causal_mask,
|
|
)
|
|
|
|
MODEL = "/tank/aimodels/gemma4-26b-a4b-it-heretic-bf16"
|
|
N = 16384
|
|
W = 1024
|
|
PAD = " " + " " * 20
|
|
|
|
cfg = AutoConfig.from_pretrained(MODEL)
|
|
text = cfg.get_text_config()
|
|
text._attn_implementation = "sdpa"
|
|
print("sliding_window %s" % text.sliding_window)
|
|
print("layers %d (%d sliding / %d full)" % (
|
|
len(text.layer_types),
|
|
text.layer_types.count("sliding_attention"),
|
|
text.layer_types.count("full_attention")))
|
|
print("_attn_implementation %s" % text._attn_implementation)
|
|
print()
|
|
|
|
|
|
def build(attn_2d, label):
|
|
batch = attn_2d.shape[0] if attn_2d is not None else 1
|
|
embeds = torch.zeros(batch, N, 8, dtype=torch.bfloat16)
|
|
pos = torch.arange(N).unsqueeze(0)
|
|
kw = dict(config=text, inputs_embeds=embeds, attention_mask=attn_2d,
|
|
past_key_values=None, position_ids=pos)
|
|
full = create_causal_mask(**kw)
|
|
slide = create_sliding_window_causal_mask(**kw)
|
|
print("--- %s ---" % label)
|
|
for name, m in (("full_attention", full), ("sliding_attention", slide)):
|
|
if m is None:
|
|
print(" %-20s None -> flash / is_causal path AVAILABLE" % name)
|
|
continue
|
|
print(" %-20s tensor shape=%s dtype=%s" % (name, tuple(m.shape), m.dtype))
|
|
allowed = m if m.dtype == torch.bool else (m == 0)
|
|
per_row = allowed[0, 0].sum(-1)
|
|
print("%sallowed/row min=%d max=%d mean=%.1f" % (
|
|
PAD, per_row.min().item(), per_row.max().item(),
|
|
per_row.float().mean().item()))
|
|
if name == "sliding_attention":
|
|
ok = per_row.max().item() <= W
|
|
print("%sBAND <= %d ? %s" % (PAD, W, "PASS" if ok else "FAIL"))
|
|
sat = (per_row >= W).nonzero()
|
|
if sat.numel():
|
|
print("%ssaturates at row %d" % (PAD, sat[0].item()))
|
|
else:
|
|
print("%slast row allows %d of %d (%s)" % (
|
|
PAD, per_row[-1].item(), N,
|
|
"causal-full OK" if per_row[-1].item() == N else "UNEXPECTED"))
|
|
print()
|
|
|
|
|
|
# 1. no 2D mask at all - the "constraints silently dropped" scenario
|
|
build(None, "attention_mask=None (no padding info)")
|
|
|
|
# 2. all-ones 2D mask - equal-length batch, no padding
|
|
build(torch.ones(2, N, dtype=torch.long), "all-ones 2D (no padding)")
|
|
|
|
# 3. REAL right-padded batch - what collate_mixed actually produces
|
|
real = torch.ones(2, N, dtype=torch.long)
|
|
real[1, 6000:] = 0
|
|
build(real, "right-padded 2D (what collate_mixed emits)")
|