Skip to content

[WIP] Sparse Attention + Recursive Weight Sharing for 16MB Efficiency #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
albertorkive wants to merge 1 commit into from
Closed

[WIP] Sparse Attention + Recursive Weight Sharing for 16MB Efficiency #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 38 additions & 0 deletions mini_run.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
#!/usr/bin/env bash
# mini_run.sh — fast local architecture testing (no CUDA needed on Mac with DEV_MODE=1)
# Usage:
# ./mini_run.sh # baseline mini-run
# WINDOW_SIZE=128 ./mini_run.sh # sliding window attention
# NUM_PHYSICAL_LAYERS=3 ./mini_run.sh # weight tying (3 blocks x 3 recurrences)
# WINDOW_SIZE=128 NUM_PHYSICAL_LAYERS=3 ./mini_run.sh # both

set -euo pipefail

ITERATIONS="${ITERATIONS:-200}"
VAL_LOSS_EVERY="${VAL_LOSS_EVERY:-50}"
SKIP_QUANT="${SKIP_QUANT:-1}"
WINDOW_SIZE="${WINDOW_SIZE:-0}"
NUM_PHYSICAL_LAYERS="${NUM_PHYSICAL_LAYERS:-0}"
DEV_MODE="${DEV_MODE:-0}"

# Detect if CUDA is available; if not, enable dev mode automatically.
python -c "import torch; exit(0 if torch.cuda.is_available() else 1)" 2>/dev/null \
|| DEV_MODE=1

echo "=== mini_run ==="
echo " ITERATIONS=$ITERATIONS VAL_LOSS_EVERY=$VAL_LOSS_EVERY"
echo " WINDOW_SIZE=$WINDOW_SIZE NUM_PHYSICAL_LAYERS=$NUM_PHYSICAL_LAYERS"
echo " SKIP_QUANT=$SKIP_QUANT DEV_MODE=$DEV_MODE"
echo "================="

ITERATIONS="$ITERATIONS" \
VAL_LOSS_EVERY="$VAL_LOSS_EVERY" \
SKIP_QUANT="$SKIP_QUANT" \
WINDOW_SIZE="$WINDOW_SIZE" \
NUM_PHYSICAL_LAYERS="$NUM_PHYSICAL_LAYERS" \
DEV_MODE="$DEV_MODE" \
python train_gpt.py 2>&1 | tee mini_run.log

echo ""
echo "=== Results ==="
grep "val_bpb" mini_run.log | tail -5
119 changes: 85 additions & 34 deletions train_gpt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,17 @@ class Hyperparameters:
rope_base = float(os.environ.get("ROPE_BASE", 10000.0))
logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0))

# Research extensions: sparse attention and recurrent weight tying.
# WINDOW_SIZE=0 → dense attention (default). WINDOW_SIZE=128 → sliding window of 128.
window_size = int(os.environ.get("WINDOW_SIZE", 0))
# NUM_PHYSICAL_LAYERS controls weight tying: if < num_layers, blocks are cycled.
# e.g. num_layers=9, num_physical_layers=3 → 3 unique blocks, each reused 3 times.
num_physical_layers = int(os.environ.get("NUM_PHYSICAL_LAYERS", 0)) # 0 = same as num_layers

# Dev/mini-run flags.
skip_quant = bool(int(os.environ.get("SKIP_QUANT", "0"))) # Skip int8+zlib for fast local runs.
dev_mode = bool(int(os.environ.get("DEV_MODE", "0"))) # Allow MPS/CPU (no CUDA required).

# Optimizer hyperparameters.
embed_lr = float(os.environ.get("EMBED_LR", 0.6))
head_lr = float(os.environ.get("HEAD_LR", 0.008))
Expand Down Expand Up @@ -255,7 +266,7 @@ def eval_val(
local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True)
x = local[:-1].reshape(-1, args.train_seq_len)
y = local[1:].reshape(-1, args.train_seq_len)
with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=True):
batch_loss = model(x, y).detach()
batch_token_count = float(y.numel())
val_loss_sum += batch_loss.to(torch.float64) * batch_token_count
Expand Down Expand Up @@ -560,6 +571,7 @@ def __init__(
num_kv_heads: int,
rope_base: float,
qk_gain_init: float,
window_size: int = 0,
):
super().__init__()
if dim % num_heads != 0:
Expand All @@ -579,6 +591,7 @@ def __init__(
self.proj._zero_init = True
self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32))
self.rotary = Rotary(self.head_dim, base=rope_base)
self.window_size = window_size

def forward(self, x: Tensor) -> Tensor:
bsz, seqlen, dim = x.shape
Expand All @@ -591,12 +604,22 @@ def forward(self, x: Tensor) -> Tensor:
q = apply_rotary_emb(q, cos, sin)
k = apply_rotary_emb(k, cos, sin)
q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None]
# Build sliding-window causal mask when window_size > 0.
# Dense causal attention (window_size=0) uses the faster is_causal=True path.
if self.window_size > 0:
idx = torch.arange(seqlen, device=x.device)
row, col = idx.unsqueeze(1), idx.unsqueeze(0)
blocked = (col > row) | (row - col >= self.window_size)
attn_mask = torch.zeros(seqlen, seqlen, device=x.device, dtype=q.dtype)
attn_mask = attn_mask.masked_fill(blocked, float("-inf"))
else:
attn_mask = None
y = F.scaled_dot_product_attention(
q,
k,
v,
attn_mask=None,
is_causal=True,
attn_mask=attn_mask,
is_causal=(attn_mask is None),
enable_gqa=(self.num_kv_heads != self.num_heads),
)
y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim)
Expand Down Expand Up @@ -626,11 +649,12 @@ def __init__(
mlp_mult: int,
rope_base: float,
qk_gain_init: float,
window_size: int = 0,
):
super().__init__()
self.attn_norm = RMSNorm()
self.mlp_norm = RMSNorm()
self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init)
self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init, window_size)
self.mlp = MLP(dim, mlp_mult)
self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32))
self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32))
Expand Down Expand Up @@ -659,10 +683,16 @@ def __init__(
logit_softcap: float,
rope_base: float,
qk_gain_init: float,
num_physical_layers: int = 0,
window_size: int = 0,
):
super().__init__()
if logit_softcap <= 0.0:
raise ValueError(f"logit_softcap must be positive, got {logit_softcap}")
# num_physical_layers < num_layers enables weight tying: blocks are cycled.
if num_physical_layers <= 0:
num_physical_layers = num_layers
self.num_physical_layers = num_physical_layers
self.tie_embeddings = tie_embeddings
self.tied_embed_init_std = tied_embed_init_std
self.logit_softcap = logit_softcap
Expand All @@ -680,8 +710,9 @@ def __init__(
mlp_mult,
rope_base,
qk_gain_init,
window_size,
)
for i in range(num_layers)
for i in range(num_physical_layers)
]
)
self.final_norm = RMSNorm()
Expand All @@ -704,13 +735,14 @@ def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor:
skips: list[Tensor] = []

# First half stores skips; second half reuses them in reverse order.
# When num_physical_layers < num_layers, blocks are cycled (weight tying).
for i in range(self.num_encoder_layers):
x = self.blocks[i](x, x0)
x = self.blocks[i % self.num_physical_layers](x, x0)
skips.append(x)
for i in range(self.num_decoder_layers):
if skips:
x = x + self.skip_weights[i].to(dtype=x.dtype)[None, None, :] * skips.pop()
x = self.blocks[self.num_encoder_layers + i](x, x0)
x = self.blocks[(self.num_encoder_layers + i) % self.num_physical_layers](x, x0)

x = self.final_norm(x).reshape(-1, x.size(-1))
targets = target_ids.reshape(-1)
Expand Down Expand Up @@ -750,23 +782,30 @@ def main() -> None:
grad_accum_steps = 8 // world_size
grad_scale = 1.0 / grad_accum_steps
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required")
device = torch.device("cuda", local_rank)
torch.cuda.set_device(device)
if not args.dev_mode:
raise RuntimeError("CUDA is required (set DEV_MODE=1 for local MPS/CPU testing)")
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
else:
device = torch.device("cuda", local_rank)
torch.cuda.set_device(device)
if distributed:
dist.init_process_group(backend="nccl", device_id=device)
dist.barrier()
master_process = rank == 0

# Fast math knobs
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp
# Fast math knobs (CUDA only)
if device.type == "cuda":
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp
enable_cudnn_sdp(False)
enable_flash_sdp(True)
enable_mem_efficient_sdp(False)
enable_math_sdp(False)

enable_cudnn_sdp(False)
enable_flash_sdp(True)
enable_mem_efficient_sdp(False)
enable_math_sdp(False)
def sync() -> None:
if device.type == "cuda":
sync()

logfile = None
if master_process:
Expand Down Expand Up @@ -800,7 +839,8 @@ def log0(msg: str, console: bool = True) -> None:
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
if device.type == "cuda":
torch.cuda.manual_seed_all(args.seed)

if not args.tokenizer_path.endswith(".model"):
raise ValueError(f"Script only setup for SentencePiece .model file: {args.tokenizer_path}")
Expand Down Expand Up @@ -835,12 +875,15 @@ def log0(msg: str, console: bool = True) -> None:
logit_softcap=args.logit_softcap,
rope_base=args.rope_base,
qk_gain_init=args.qk_gain_init,
num_physical_layers=args.num_physical_layers,
window_size=args.window_size,
).to(device).bfloat16()
for module in base_model.modules():
if isinstance(module, CastedLinear):
module.float()
restore_low_dim_params_to_fp32(base_model)
compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True)
# Skip torch.compile in dev mode (MPS/CPU) to avoid backend issues.
compiled_model = base_model if args.dev_mode else torch.compile(base_model, dynamic=False, fullgraph=True)
model: nn.Module = DDP(compiled_model, device_ids=[local_rank], broadcast_buffers=False) if distributed else compiled_model

# Optimizer split:
Expand All @@ -866,7 +909,7 @@ def log0(msg: str, console: bool = True) -> None:
[{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}],
betas=(args.beta1, args.beta2),
eps=args.adam_eps,
fused=True,
fused=(device.type == "cuda"),
)
optimizer_muon = Muon(
matrix_params,
Expand All @@ -880,15 +923,15 @@ def log0(msg: str, console: bool = True) -> None:
[{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}],
betas=(args.beta1, args.beta2),
eps=args.adam_eps,
fused=True,
fused=(device.type == "cuda"),
)
optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar]
if base_model.lm_head is not None:
optimizer_head = torch.optim.Adam(
[{"params": [base_model.lm_head.weight], "lr": args.head_lr, "base_lr": args.head_lr}],
betas=(args.beta1, args.beta2),
eps=args.adam_eps,
fused=True,
fused=(device.type == "cuda"),
)
optimizers.insert(1, optimizer_head)

Expand Down Expand Up @@ -944,7 +987,7 @@ def lr_mul(step: int, elapsed_ms: float) -> float:
if distributed:
model.require_backward_grad_sync = micro_step == grad_accum_steps - 1
x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps)
with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=True):
warmup_loss = model(x, y)
(warmup_loss * grad_scale).backward()
for opt in optimizers:
Expand All @@ -966,7 +1009,7 @@ def lr_mul(step: int, elapsed_ms: float) -> float:

training_time_ms = 0.0
stop_after_step: int | None = None
torch.cuda.synchronize()
sync()
t0 = time.perf_counter()

step = 0
Expand All @@ -975,7 +1018,7 @@ def lr_mul(step: int, elapsed_ms: float) -> float:

should_validate = last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0)
if should_validate:
torch.cuda.synchronize()
sync()
training_time_ms += 1000.0 * (time.perf_counter() - t0)
val_loss, val_bpb = eval_val(
args,
Expand All @@ -993,7 +1036,7 @@ def lr_mul(step: int, elapsed_ms: float) -> float:
f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} "
f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms"
)
torch.cuda.synchronize()
sync()
t0 = time.perf_counter()

if last_step:
Expand All @@ -1012,7 +1055,7 @@ def lr_mul(step: int, elapsed_ms: float) -> float:
if distributed:
model.require_backward_grad_sync = micro_step == grad_accum_steps - 1
x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps)
with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=True):
loss = model(x, y)
train_loss += loss.detach()
(loss * grad_scale).backward()
Expand Down Expand Up @@ -1054,16 +1097,24 @@ def lr_mul(step: int, elapsed_ms: float) -> float:
if stop_after_step is None and reached_cap:
stop_after_step = step

log0(
f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB "
f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB"
)
if device.type == "cuda":
log0(
f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB "
f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB"
)

# -----------------------------
# SERIALIZATION + ROUNDTRIP VALIDATION
# -----------------------------
# Save the raw state (useful for debugging/loading in PyTorch directly), then always produce
# the compressed int8+zlib artifact and validate the round-tripped weights.
# SKIP_QUANT=1 skips this section entirely (fast local dev / mini-runs).
if args.skip_quant:
if master_process:
log0("skip_quant=True: skipping serialization and roundtrip validation")
if distributed:
dist.destroy_process_group()
return

if master_process:
torch.save(base_model.state_dict(), "final_model.pt")
Expand Down Expand Up @@ -1097,7 +1148,7 @@ def lr_mul(step: int, elapsed_ms: float) -> float:
quant_blob_disk = f.read()
quant_state = torch.load(io.BytesIO(zlib.decompress(quant_blob_disk)), map_location="cpu")
base_model.load_state_dict(dequantize_state_dict_int8(quant_state), strict=True)
torch.cuda.synchronize()
sync()
t_qeval = time.perf_counter()
q_val_loss, q_val_bpb = eval_val(
args,
Expand All @@ -1111,7 +1162,7 @@ def lr_mul(step: int, elapsed_ms: float) -> float:
has_leading_space_lut,
is_boundary_token_lut,
)
torch.cuda.synchronize()
sync()
log0(
f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} "
f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms"
Expand Down