Skip to content

feat: depth recurrence + cosine recovery TTT#697

Open
Danishlynx wants to merge 2 commits intoopenai:mainfrom
Danishlynx:feat/combined-best
Open

feat: depth recurrence + cosine recovery TTT#697
Danishlynx wants to merge 2 commits intoopenai:mainfrom
Danishlynx:feat/combined-best

Conversation

@Danishlynx
Copy link

Based on merged SOTA (PR #549 stack + LeakyReLU² + Legal TTT, 1.1194 bpb):

  1. Depth recurrence: repeat layers 4-5 → 13 virtual layers from 11 physical

    • Per-repetition learnable scale parameters
    • U-Net skip connections adapted for virtual layer count
    • DEPTH_RECURRENCE=4,5 env var
  2. Enhanced TTT with cosine recovery phase:

    • After standard score-first TTT, runs N additional cosine-LR epochs on all scored data to repair int6 quantization damage
    • Re-scores with standard sliding window eval
    • TTT_RECOVERY_EPOCHS=20, TTT_RECOVERY_LR=0.001 env vars
  3. FlashAttention 3 fallback to SDPA for non-Hopper GPUs

    • Manual GQA head repeat for PyTorch <2.5 compatibility

Smoke-tested on 1xH100 SXM 80GB. Both features validated.

Danishlynx and others added 2 commits March 25, 2026 14:21
Based on merged SOTA (PR openai#549 stack + LeakyReLU² + Legal TTT, 1.1194 bpb):

1. Depth recurrence: repeat layers 4-5 → 13 virtual layers from 11 physical
   - Per-repetition learnable scale parameters
   - U-Net skip connections adapted for virtual layer count
   - DEPTH_RECURRENCE=4,5 env var

2. Enhanced TTT with cosine recovery phase:
   - After standard score-first TTT, runs N additional cosine-LR epochs
     on all scored data to repair int6 quantization damage
   - Re-scores with standard sliding window eval
   - TTT_RECOVERY_EPOCHS=20, TTT_RECOVERY_LR=0.001 env vars

3. FlashAttention 3 fallback to SDPA for non-Hopper GPUs
   - Manual GQA head repeat for PyTorch <2.5 compatibility

Smoke-tested on 1xH100 SXM 80GB. Both features validated.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- fullgraph=True → fullgraph=False for torch.compile (conditional
  branches in _run_layers break fullgraph)
- Create fresh uncompiled model for TTT eval to avoid stale inference
  tensor state from compiled eval model
- Clear Rotary cos/sin caches when transitioning between
  inference_mode (scoring) and train mode (adaptation) to prevent
  "Inference tensors cannot be saved for backward" errors
- Manual GQA head repeat for PyTorch <2.5 SDPA compatibility

Validated: TTT now runs end-to-end on 1xH100, achieving 1.3859 bpb
(from 1.5158 post-quant baseline, -0.13 bpb improvement)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant