-
Notifications
You must be signed in to change notification settings - Fork 3.2k
[Record] 11L Depth Recurrence + EMA Tuning (0.9965) — val_bpb 1.0925 #1421
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.
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
Open
X-Abhishek-X
wants to merge
3
commits into
openai:main
Choose a base branch
from
X-Abhishek-X:record/11L-depth-recurrence-ema-0.9965
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
3 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
62 changes: 62 additions & 0 deletions
62
records/track_10min_16mb/2026-04-06_11L_DepthRecurrence_EMA0.9965_1.0925/README.md
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,62 @@ | ||
| ## Record: 11L Depth Recurrence + EMA Tuning (0.9965) (val_bpb: 1.0925) | ||
|
|
||
| **val_bpb: 1.0925** (sliding window stride=64, 3-seed mean) | **15.95 MB** (mean) | 8xH100 SXM, 590s | ||
|
|
||
| ### Key Change | ||
|
|
||
| EMA decay hyperparameter refinement on top of PR #1334's (@aryanbhosale) depth recurrence architecture: | ||
|
|
||
| | Parameter | Baseline | This | Impact | | ||
| |-----------|----------|------|--------| | ||
| | **EMA decay** | 0.997 | 0.9965 | Stabilized post-quantization performance, reduced selective pruning to ~290K values | | ||
|
|
||
| ### EMA Decay Tuning | ||
|
|
||
| By lowering the EMA decay from 0.997 to 0.9965, the exponential moving average assigns slightly more weight to recent training steps. This produces a final checkpoint that quantizes more cleanly under GPTQ int6, reducing the number of values requiring selective pruning (~290K vs baseline). | ||
|
|
||
| ### Results (3 seeds, 8xH100 SXM) | ||
|
|
||
| | Seed | Pre-quant BPB | Sliding BPB (s64) | Artifact | | ||
| |------|---------------|-------------------|----------| | ||
| | 42 | 1.0965 | **1.0921** | 15,954,858 B | | ||
| | 1337 | 1.0973 | **1.0928** | 15,959,674 B | | ||
| | 2024 | 1.0969 | **1.0926** | 15,948,766 B | | ||
|
|
||
| **Mean: 1.0925 | Std: 0.0004** | All artifacts under 16,000,000 bytes | ||
|
|
||
| ### Architecture (from PR #1334) | ||
|
|
||
| - 11 transformer layers, 512-dim, 8 heads (4 KV heads, GQA) | ||
| - Depth recurrence: layers 4,5 repeat (virtual 13 layers), activated at step 3000 | ||
| - Skip gates (learnable residual gating) | ||
| - Shared Value Embedding (dim=128, layers 9,10) | ||
| - Tied embeddings, logit softcap=30.0 | ||
| - SP4096 tokenizer (SentencePiece BPE) | ||
|
|
||
| ### Training | ||
|
|
||
| - FlashAttention 3 (Hopper-optimized) | ||
| - Muon optimizer (matrices): lr=0.02, momentum=0.99, WD=0.09, backend_steps=5 | ||
| - Adam (head params): lr=0.008, fused=True | ||
| - AdamW (embeddings): lr=0.6, WD=0.09, fused=True | ||
| - AdamW (scalars): lr=0.02, WD=0.02, fused=True | ||
| - Gradient clip: 0.3 | ||
| - Batch: 786,432 tokens/step, seq_len=2048 | ||
| - Warmdown: 66.7% of training | ||
| - **EMA**: decay=0.9965, every step | ||
| - Wallclock cap: 600s (590s effective, 10s reserved for GPTQ) | ||
|
|
||
| ### Quantization | ||
|
|
||
| - GPTQ int6 with percdamp=0.05, 64 calibration batches | ||
| - Selective pruning of lowest-error values to fit 16MB | ||
| - Brotli compression | ||
| - ~290K values pruned (minimal impact) | ||
|
|
||
| ### Reproducibility | ||
|
|
||
| All 3 seeds produce valid artifacts under 16MB with tight variance (std=0.0004 BPB). Training completes in ~590s with ~5200-5400 steps depending on seed. | ||
|
|
||
| ### Attribution | ||
|
|
||
| Base architecture and training recipe from PR #1334 by @aryanbhosale. | ||
10 changes: 10 additions & 0 deletions
10
records/track_10min_16mb/2026-04-06_11L_DepthRecurrence_EMA0.9965_1.0925/submission.json
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,10 @@ | ||
| { | ||
| "author": "Abhishek Leji", | ||
| "github_id": "X-Abhishek-X", | ||
| "name": "Record: 11L Depth Recurrence + EMA Tuning (0.9965)", | ||
| "blurb": "EMA decay tuned to 0.9965 for stabilized post-quantization performance, built on PR #1334 (aryanbhosale) depth recurrence architecture (11L, skip gates, VE128, GPTQ int6+brotli, sliding window eval).", | ||
| "date": "2026-04-06T00:00:00Z", | ||
| "val_loss": 2.51370050, | ||
| "val_bpb": 1.09247668, | ||
| "bytes_total": 15954433 | ||
| } |
133 changes: 133 additions & 0 deletions
133
records/track_10min_16mb/2026-04-06_11L_DepthRecurrence_EMA0.9965_1.0925/train.log
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,133 @@ | ||
| W0406 15:28:12.622000 47154 torch/distributed/run.py:803] | ||
| W0406 15:28:12.622000 47154 torch/distributed/run.py:803] ***************************************** | ||
| W0406 15:28:12.622000 47154 torch/distributed/run.py:803] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. | ||
| W0406 15:28:12.622000 47154 torch/distributed/run.py:803] ***************************************** | ||
| Hyperparameters: | ||
| adam_eps: 1e-08 | ||
| adam_wd: 0.02 | ||
| beta1: 0.9 | ||
| beta2: 0.95 | ||
| compressor: brotli | ||
| data_dir: ./data/ | ||
| datasets_dir: ./data/datasets/fineweb10B_sp4096 | ||
| distributed: True | ||
| ema_decay: 0.9965 | ||
| embed_lr: 0.6 | ||
| embed_wd: 0.09 | ||
| embedding_dim: 512 | ||
| eval_seq_len: 2048 | ||
| eval_stride: 64 | ||
| gptq_calibration_batches: 64 | ||
| gptq_enabled: True | ||
| gptq_reserve_seconds: 10.0 | ||
| grad_accum_steps: 1 | ||
| grad_clip_norm: 0.3 | ||
| head_lr: 0.008 | ||
| is_main_process: True | ||
| iterations: 20000 | ||
| ln_scale: True | ||
| local_rank: 0 | ||
| logfile: logs/21b90172-c576-4e07-bf1d-b88115ce3546.txt | ||
| logit_softcap: 30.0 | ||
| matrix_lr: 0.02 | ||
| max_wallclock_seconds: 600.0 | ||
| min_lr: 0.0 | ||
| mlp_mult: 4.0 | ||
| model_dim: 512 | ||
| model_path: final_model.pt | ||
| muon_backend_steps: 5 | ||
| muon_beta2: 0.95 | ||
| muon_momentum: 0.99 | ||
| muon_momentum_warmup_start: 0.92 | ||
| muon_momentum_warmup_steps: 1500 | ||
| muon_wd: 0.09 | ||
| num_heads: 8 | ||
| num_kv_heads: 4 | ||
| num_layers: 11 | ||
| parallel_start_layer: 7 | ||
| qk_gain_init: 5.0 | ||
| quantized_model_path: final_model.int6.ptz | ||
| rank: 0 | ||
| recur_layers: 4,5 | ||
| recur_start_step: 3000 | ||
| rope_base: 10000.0 | ||
| rope_dims: 16 | ||
| rope_train_seq_len: 2048 | ||
| run_id: 21b90172-c576-4e07-bf1d-b88115ce3546 | ||
| scalar_lr: 0.02 | ||
| seed: 42 | ||
| skip_gates_enabled: True | ||
| sliding_window_enabled: True | ||
| tie_embeddings: True | ||
| tied_embed_init_std: 0.005 | ||
| tied_embed_lr: 0.03 | ||
| tokenizer_path: ./data/tokenizers/fineweb_4096_bpe.model | ||
| train_batch_tokens: 786432 | ||
| train_files: ./data/datasets/fineweb10B_sp4096/fineweb_train_*.bin | ||
| train_log_every: 500 | ||
| train_seq_len: 2048 | ||
| ttt_batch_seqs: 32 | ||
| ttt_chunk_tokens: 32768 | ||
| ttt_enabled: False | ||
| ttt_epochs: 3 | ||
| ttt_freeze_blocks: 0 | ||
| ttt_grad_clip: 1.0 | ||
| ttt_lr: 0.002 | ||
| ttt_momentum: 0.9 | ||
| val_batch_tokens: 524288 | ||
| val_files: ./data/datasets/fineweb10B_sp4096/fineweb_val_*.bin | ||
| val_loss_every: 4000 | ||
| ve_dim: 128 | ||
| ve_enabled: True | ||
| ve_layers: 9,10 | ||
| vocab_size: 4096 | ||
| warmdown_frac: 0.667 | ||
| warmup_steps: 20 | ||
| world_size: 8 | ||
| xsa_last_n: 11 | ||
| train_shards: 143 | ||
| val_tokens: 45508608 | ||
| model_params:34401372 | ||
| gptq:reserving 10s, effective=590000ms | ||
| warmup_step: 1/20 | ||
| warmup_step: 2/20 | ||
| warmup_step: 3/20 | ||
| warmup_step: 4/20 | ||
| warmup_step: 5/20 | ||
| warmup_step: 6/20 | ||
| warmup_step: 10/20 | ||
| warmup_step: 20/20 | ||
| 0/20000 val_loss: 8.3187 val_bpb: 3.6152 | ||
| 1/20000 train_loss: 8.3178 train_time: 0.0m tok/s: 8459847 | ||
| 2/20000 train_loss: 12.0900 train_time: 0.0m tok/s: 8353619 | ||
| 3/20000 train_loss: 10.6985 train_time: 0.0m tok/s: 8256976 | ||
| 4/20000 train_loss: 8.9900 train_time: 0.0m tok/s: 8042542 | ||
| 5/20000 train_loss: 7.7467 train_time: 0.0m tok/s: 8037005 | ||
| 500/20000 train_loss: 2.9900 train_time: 0.8m tok/s: 7915080 | ||
| 1000/20000 train_loss: 2.9924 train_time: 1.7m tok/s: 7884266 | ||
| 1500/20000 train_loss: 2.9051 train_time: 2.5m tok/s: 7879749 | ||
| 2000/20000 train_loss: 2.7520 train_time: 3.3m tok/s: 7875834 | ||
| 2500/20000 train_loss: 2.7534 train_time: 4.2m tok/s: 7873543 | ||
| 3000/20000 train_loss: 2.7280 train_time: 5.0m tok/s: 7873311 | ||
| recurrence:activated at step 3000, virtual_layers=[0, 1, 2, 3, 4, 5, 4, 5, 6, 7, 8, 9, 10] | ||
| 3500/20000 train_loss: 2.6339 train_time: 6.1m tok/s: 7465954 | ||
| 4000/20000 train_loss: 2.6341 train_time: 7.1m tok/s: 7375160 | ||
| 4000/20000 val_loss: 2.6390 val_bpb: 1.1469 | ||
| 4500/20000 train_loss: 2.5800 train_time: 8.1m tok/s: 7306413 | ||
| 5000/20000 train_loss: 2.5396 train_time: 9.0m tok/s: 7251921 | ||
| 5413/20000 val_loss: 2.5260 val_bpb: 1.0978 | ||
| stopping_early: wallclock_cap train_time: 590010ms step: 5413/20000 | ||
| peak memory allocated: 30120 MiB reserved: 30154 MiB | ||
| ema:applying EMA weights | ||
| pre-quantization post-ema val_loss:2.52306664 val_bpb:1.09649570 eval_time:2024ms | ||
| Serialized model: 132406149 bytes | ||
| Code size: 83566 bytes | ||
| GPTQ:collecting Hessians from calibration data... | ||
| GPTQ:collected 66 Hessians in 9.8s | ||
| GPTQ quantization: 66 layers with full GPTQ, 0 fallback to clip-search | ||
| selective_prune: unpruned=16.04MB target=16.0MB | ||
| selective_prune: pruning 290456/9399054 lowest-error ±1 values (excess=36307B) | ||
| Serialized model int6+brotli: 15871292 bytes | ||
| Total submission size int6+brotli: 15954858 bytes | ||
| final_int6_roundtrip val_loss:2.55549554 val_bpb:1.11058892 eval_time:8386ms | ||
| final_int6_sliding_window val_loss:2.51297660 val_bpb:1.09211068 eval_time:76584ms |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The README labels this as a “Record” and frames it as an improvement over PR #1334, but the PR metadata you reference lists PR #1334 with a lower (better)
val_bpb(1.0897). Please clarify the baseline/track comparison or adjust the wording so the record claim is unambiguous and consistent with the referenced results.