第59日: バックグラウンド実行、メトリクステレメトリと動的解析

# 1. Start the run detached. tee -a appends, so a restart never truncates earlier steps.
tmux new-session -d -s pretrain 'python3 train.py 2>&1 | tee -a training.log'

# 2. Confirm that the session and the training process both exist.
tmux ls # expect: pretrain: 1 windows (created ...)
ps aux | grep '[t]rain.py' # the [t] pattern stops grep from matching itself

# 3. Record two log tails a minute apart; the step number must grow.
tail -n 2 training.log; sleep 60; tail -n 2 training.log

# 4. Disconnect test: close the SSH client (do not type exit inside tmux), reconnect, then repeat
# step 2 and step 3. The process survives because the tmux server, not your SSH login shell,
# is its parent. nohup python3 train.py >> training.log 2>&1 & is the no-tmux fallback.
tmux attach -t pretrain # watch live; Ctrl-B then d detaches again without stopping training

# 5. Stop on purpose at the end. Ctrl-C sends SIGINT; train.py should save a checkpoint and exit.
tmux send-keys -t pretrain C-c

# Evidence to keep: the two tails (step grew), the ps line before and after the disconnect,
# and the wall-clock hours. Only the minutes you spent reading telemetry count as study time.

秒間トークン数 = 32,768 / 0.25 = 131,072 トークン/秒。

トークンあたり FLOPs = 6 * N。正確な数 N = 124,356,864 では 746,141,184 FLOPs/トークン。
達成 FLOPs/秒 = 746,141,184 * 131,072 = 9.78e13 FLOPs/秒 = 97.80 TFLOP/s。
MFU = 97.80 / 312.0 = 31.35%。
丸めた N = 124e6 では 31.26%、Day 47 の 124,439,808 では 31.37%。N の違いは MFU の小数第 2 位しか動かさないので、どの N を使ったかを明記する。

ピーク値：312 TFLOP/s は A100 の密な bf16（および fp16）テンソルコアのピークである。fp32 のピーク（19.5 TFLOP/s）でも、スパース時の値（624）でもない。MFU は訓練に使う dtype のピークに対してだけ意味を持つ。

(a) 6N は重み行列の積だけを数える。注意スコアの積は訓練（順伝播と逆伝播）でトークンあたり約 12 * n_layer * T * d_model = 12 * 12 * 1024 * 768 = 113,246,208 FLOPs を加える。これを含めると (746,141,184 + 113,246,208) * 131,072 / 312e12 = 36.10%。どちらの慣例を使うかを報告する。一般的で、このコースが使うのは 6N である。

(b) 崩壊：V 個のトークンを一様に推測するモデルの損失は ln(V) = ln(50257) ≈ 10.825。約 10.8 から下がる損失は学習を意味し、10.8 付近に戻ってとどまる損失は崩壊を意味する。NaN や急上昇は発散である。

結果の読み方：1 台の A100 上の小さなモデルでは 31% の MFU はよくある値である。MFU が低い原因は多くの場合、データローダーの停滞、小さなバッチ、融合されていないカーネルであり、プロファイルする理由であって訓練の誤りではない。

import os
import torch
import torch.nn as nn

def make():
torch.manual_seed(0)
model = nn.Linear(8, 1)
return model, torch.optim.AdamW(model.parameters(), lr=1e-2)

def batch(step):
# The batch depends only on the step number, so a resumed run sees the same data.
g = torch.Generator().manual_seed(step)
x = torch.randn(16, 8, generator=g)
return x, x.sum(dim=1, keepdim=True)

def save_atomic(state, path):
tmp = path + '.tmp'
torch.save(state, tmp)
os.replace(tmp, path) # atomic rename: a reader never sees half a file (Day 55)

def train(model, opt, start, stop, ckpt=None, save_every=None):
losses = {}
for step in range(start, stop + 1):
x, y = batch(step)
loss = nn.functional.mse_loss(model(x), y)
opt.zero_grad()
loss.backward()
opt.step()
losses[step] = loss.item()
if save_every and step % save_every == 0:
# Save model AND optimizer state; AdamW moments are part of the training state.
save_atomic({'step': step, 'loss': losses[step],
'model': model.state_dict(), 'opt': opt.state_dict()}, ckpt)
return losses

# Baseline: one uninterrupted run from step 1 to 110.
baseline = train(*make(), 1, 110)

# Interrupted run: it saves every 50 steps. Stopping after step 100 stands in for a
# SIGTERM handler that finishes the current step, saves, and exits.
train(*make(), 1, 100, ckpt='ckpt.pt', save_every=50)

# Resume in a fresh process state: rebuild, load, continue from the saved step + 1.
model, opt = make()
state = torch.load('ckpt.pt')
model.load_state_dict(state['model'])
opt.load_state_dict(state['opt'])
resumed = train(model, opt, state['step'] + 1, 110)

resumed_step, last_saved_step = min(resumed) - 1, state['step']
assert resumed_step == last_saved_step == 100
assert abs(state['loss'] - baseline[100]) < 1e-5 # step-100 loss survived the restart
assert all(abs(resumed[s] - baseline[s]) < 1e-5 for s in range(101, 111))
print('resume verified: steps 101-110 match the uninterrupted run')

# Failure you should provoke once: resume WITHOUT opt.load_state_dict. The step-101 loss still
# matches (same weights), but later steps drift because AdamW restarts with zero moments.
# Retention: keep only the last 3 checkpoints, e.g. ckpt-0050.pt, ckpt-0100.pt, ckpt-0150.pt,
# and delete the oldest only after the newest os.replace has succeeded.

import math
import re
import time
from collections import deque

LINE = re.compile(r'step: (\d+), loss: ([^,]+), dt: ([0-9.eE+-]+)')

class TelemetryMonitor:
def __init__(self, model_params, tokens_per_step, peak_flops, window=10, max_loss=20.0):
self.params = model_params
self.tokens_per_step = tokens_per_step
self.peak_flops = peak_flops
self.max_loss = max_loss
self.dts = deque(maxlen=window) # rolling window for tokens/sec
self.losses = deque(maxlen=window) # rolling window for the loss moving average
self.mfus = [] # every step, for the run-average MFU

def parse(self, line):
match = LINE.match(line.strip())
if not match:
return None # ignore banners, warnings and partial lines
return int(match[1]), float(match[2]), float(match[3]) # float('nan') parses 'nan'

def process_step(self, step, loss, dt):
if not math.isfinite(loss):
return {'status': 'CRITICAL', 'step': step, 'reason': 'non-finite loss', 'action': 'terminate now'}
if loss > self.max_loss:
return {'status': 'CRITICAL', 'step': step, 'reason': f'loss {loss} > {self.max_loss}', 'action': 'terminate now'}
mfu = 6 * self.params * self.tokens_per_step / (dt * self.peak_flops)
self.dts.append(dt)
self.losses.append(loss)
self.mfus.append(mfu)
return {'status': 'OK', 'step': step, 'mfu': mfu,
'avg_mfu': sum(self.mfus) / len(self.mfus),
'tokens_per_sec': self.tokens_per_step * len(self.dts) / sum(self.dts),
'loss_ma': sum(self.losses) / len(self.losses)}

def follow(self, log_path, poll_seconds=1.0):
# Non-blocking tail: start at the end of the file and yield None when nothing is new,
# so the caller can also check wall-clock limits between lines.
with open(log_path) as f:
f.seek(0, 2)
while True:
line = f.readline()
if not line:
yield None
time.sleep(poll_seconds)
continue
parsed = self.parse(line)
if parsed:
yield self.process_step(*parsed)

# Simulated stream: 20 healthy steps, then an injected NaN.
monitor = TelemetryMonitor(124_356_864, 32_768, 312e12)
stream = [f'step: {s}, loss: {4.0 - 0.05 * s:.3f}, dt: 0.25' for s in range(1, 21)]
stream.append('step: 21, loss: nan, dt: 0.25')
reports = [monitor.process_step(*monitor.parse(line)) for line in stream]

assert all(r['status'] == 'OK' for r in reports[:20])
assert abs(reports[19]['avg_mfu'] - 0.3135) < 1e-4 # 31.35%, the Activity 2 value
assert reports[19]['tokens_per_sec'] == 131072.0
assert reports[20]['status'] == 'CRITICAL' and reports[20]['step'] == 21 # flagged on the NaN line itself
assert monitor.process_step(22, 25.0, 0.25)['reason'].startswith('loss 25.0') # divergence rule
print(reports[19]['avg_mfu'], reports[20])