WIP: sparse attention + recursive weight sharing skeleton

Adds three research levers to train_gpt.py (1177 lines, under 1500 cap):

- Sliding window attention (WINDOW_SIZE env var, default 0 = dense)
  Causal attention restricted to a local window; reduces compute and
  may allow wider/deeper models within the 16MB artifact budget.

- Recursive weight tying (NUM_PHYSICAL_LAYERS env var)
  N physical blocks reused across NUM_LAYERS logical layers. E.g.
  NUM_PHYSICAL_LAYERS=3 with NUM_LAYERS=9 gives 3x effective depth
  at the same parameter count.

- Dev/mini-run support (DEV_MODE, SKIP_QUANT env vars)
  DEV_MODE=1 allows MPS/CPU fallback for local architecture testing.
  SKIP_QUANT=1 skips int8+zlib serialization for fast iteration.

Also adds mini_run.sh convenience script for quick local runs.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: albertorkive

mini_run.sh

@@ -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

train_gpt.py

@@ -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))
@@ -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
@@ -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:
@@ -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
@@ -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)
@@ -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))
@@ -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
@@ -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()
@@ -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)
@@ -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:
@@ -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}")
@@ -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:
@@ -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,
@@ -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)
 
@@ -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:
@@ -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
@@ -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,
@@ -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:
@@ -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()
@@ -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")
@@ -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,
@@ -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"