Add 3d optimizer sweep - #894
Open
klei22 wants to merge 5 commits into
Open
Conversation
There was a problem hiding this comment.
Pull request overview
Adds a richer “3D trajectory” sweep workflow for the digits-3d demo, expanding sweep dimensions (optimizer × radius × token-availability schedule) and adding tooling to package completed runs as a GitHub Pages–friendly static site.
Changes:
- Extends digits-3d sweep generation to cover optimizer comparisons, radius modes, and drop/add/duty-cycle schedules, and exports the extra metadata into run JSON + manifest.
- Updates the Three.js sweep selector/viewer UI and dataset generator to support dropped-symbol schedules while keeping vocabulary stable.
- Adds a packaging script + Python utility to build a deployable static site from completed runs.
Reviewed changes
Copilot reviewed 16 out of 17 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| train.py | Adjusts Muon parameter routing when --muon_include_all_weights is enabled. |
| train_args.py | Updates CLI help text for Muon “include all weights” behavior. |
| analysis/export_3d_token_trajectories.py | Exports schedule + optimizer metadata into trajectory JSON payloads. |
| analysis/update_3d_sweep_manifest.py | Adds schedule/optimizer fields into sweep manifest entries. |
| analysis/package_3d_trajectory_site.py | New: packages completed runs + viewer + selector as a static site. |
| data/digits_3d/prepare.py | Adds --dropout-count to generate splits that exclude trailing trained symbols while retaining vocab. |
| report/threejs/digits-3d/index.html | Updates viewer legend + rendering logic for excluded segments and optimizer metadata display. |
| report/threejs/digits-3d/sweep.html | Adds new filters (radius/schedule/optimizer/etc.) and richer run cards. |
| report/threejs/digits-3d/README.md | Documents new sweep dimensions and GitHub Pages packaging flow. |
| demos/digits_3d_trajectory_demo.sh | Adds schedule execution (drop/add/duty-cycle) and optimizer selection for demos. |
| demos/digits_3d_trajectory_sweep.sh | Expands sweep loops across new dimensions and regenerates manifest per run. |
| demos/package_digits_3d_github_pages.sh | New: wrapper to build (and optionally publish) a Pages-ready directory. |
| tests/test_digits_3d_dataset.py | Adds coverage for dropped-symbol dataset behavior. |
| tests/test_update_3d_sweep_manifest.py | Extends manifest test to assert new fields. |
| tests/test_package_3d_trajectory_site.py | New: tests static-site packaging behavior. |
| .gitignore | Ignores generated dist/digits-3d-site/ output. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| # Newton-Schulz orthogonalization is defined for matrices. | ||
| # "All weights" includes embeddings/heads, but leaves | ||
| # scalar/vector parameters on Muon's auxiliary Adam path. | ||
| return param.ndim >= muon_min_ndim |
| "wte_weight_tying": payload.get("wte_weight_tying", True), | ||
| "dropped_tokens": len(payload.get("dropped_tokens", [])), | ||
| "dropout_iteration": payload.get("dropout_iteration"), | ||
| "final_iteration": payload.get("frames", [{}])[-1].get("iteration"), |
Comment on lines
10
to
+14
| function render(){ | ||
| const filtered=state.runs.filter(run=>(!dimSelect.value||run.embedding_dim==dimSelect.value)&&(!trainedSelect.value||run.trained_tokens==trainedSelect.value)&&(!heldSelect.value||run.held_out_tokens==heldSelect.value)&&(!tyingSelect.value||String(run.wte_weight_tying)===tyingSelect.value)); | ||
| const valueOf=run=>run.transition_mode==='duty_cycle'?run.duty_cycle:Math.round(100*run.dropout_iteration/run.final_iteration),radiusOf=run=>run.fixed_norm==null?'free':String(run.fixed_norm);const filtered=state.runs.filter(run=>(!dimSelect.value||run.embedding_dim==dimSelect.value)&&(!radiusSelect.value||radiusOf(run)===radiusSelect.value)&&(!trainedSelect.value||run.trained_tokens==trainedSelect.value)&&(!heldSelect.value||run.held_out_tokens==heldSelect.value)&&(!droppedSelect.value||run.dropped_tokens==droppedSelect.value)&&(!scheduleSelect.value||run.transition_mode===scheduleSelect.value)&&(!dropPercentSelect.value||valueOf(run)==dropPercentSelect.value)&&(!optimizerSelect.value||run.optimizer?.name===optimizerSelect.value)&&(!tyingSelect.value||String(run.wte_weight_tying)===tyingSelect.value)); | ||
| runList.innerHTML=''; | ||
| if(!filtered.includes(state.selected))state.selected=null; | ||
| filtered.forEach(run=>{const button=document.createElement('button');button.textContent=`${run.embedding_dim}D · ${run.trained_tokens} trained · ${run.held_out_tokens} held · ${run.wte_weight_tying?'tied':'untied'} · ${run.fixed_norm==null?'free':`radius ${Number(run.fixed_norm).toFixed(2)}`}`;button.className=run===state.selected?'active':'';button.onclick=()=>selectRun(run);runList.appendChild(button)}); | ||
| filtered.forEach(run=>{const schedule=run.transition_mode==='duty_cycle'?`${run.duty_cycle}% duty`:`${run.transition_mode} @ ${Math.round(100*run.dropout_iteration/run.final_iteration)}%`,optimizer=run.optimizer?.name||'legacy';const button=document.createElement('button');button.textContent=`${run.embedding_dim}D · ${optimizer}${run.optimizer?.muon_include_all_weights?' (full)':''} · wd ${run.optimizer?.weight_decay??'—'} · ${run.trained_tokens} trained · ${run.held_out_tokens} held · ${run.dropped_tokens} affected · ${schedule} · ${run.wte_weight_tying?'tied':'untied'} · ${run.fixed_norm==null?'free':`radius ${Number(run.fixed_norm).toFixed(2)}`}`;button.className=run===state.selected?'active':'';button.onclick=()=>selectRun(run);runList.appendChild(button)}); |
|
|
||
| Example result: | ||
| http://localhost:8000/report/threejs/digits-3d/index.html?data=runs/dim-3_digits-10_letters-10_sqrt_dim_tied.json | ||
| http://localhost:8000/report/threejs/digits-3d/index.html?data=runs/dim-3_digits-10_letters-10_sqrt_dim_tied_adam_drop-1_drop-at-40pct.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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
No description provided.