Skip to content

Improve implementations of Dion2 #21

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

Open
kwangjunahn wants to merge 14 commits into
base: main
Choose a base branch
from
Open

Improve implementations of Dion2 #21

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
c25daf5
dion2 improvements
kwangjunahn Dec 16, 2025
cb7f0ab
some tests
kwangjunahn Dec 16, 2025
ed196d3
debug stmt
kwangjunahn Dec 17, 2025
4d22079
clean up
kwangjunahn Dec 17, 2025
cf0e5fe
improved dion2 impl
kwangjunahn Dec 18, 2025
843e4c3
train.py update, readme update
kwangjunahn Dec 18, 2025
3f8a1b3
readme
kwangjunahn Dec 18, 2025
c54135f
readme
kwangjunahn Dec 18, 2025
2dc8032
read me
kwangjunahn Dec 18, 2025
8ff2aad
readme update
kwangjunahn Dec 18, 2025
aee7251
readme reorganize
kwangjunahn Dec 18, 2025
d3eeab3
readme clean up
kwangjunahn Dec 18, 2025
1b61ff5
black format
kwangjunahn Dec 18, 2025
df933c3
comment about L1 norm
kwangjunahn Dec 23, 2025
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
3 changes: 2 additions & 1 deletion .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,4 +7,5 @@ wandb/
output/
.venv/
submit.ipynb
aztool/
aztool/
dion.egg-info/
242 changes: 126 additions & 116 deletions README.md
View file
Open in desktop

Large diffs are not rendered by default.

21 changes: 9 additions & 12 deletions configs/dion2_160m.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,32 +14,29 @@ num_iterations: 3000
warmup_ratio: 0.0
warmdown_ratio: 0.2

# — Optimizer & Hyperparameters —
mu: 0.95
weight_decay: 0.01

# — Validation & Checkpointing —
val_loss_every: 125
val_tokens: 10485760
save_every: 0
val_tokens: 10485760

# — Weights & Biases logging —
wandb_project_name: gpt-train
wandb_job_name: null
no_wandb: false

# — Distributed training —
dp_size: null # data‐parallel size
fs_size: null # FSDP size
tp_size: null # DO NOT USE TP for Dion2

# — Miscellaneous flags —
debug: false
no_compile: false
no_triton: false

# — Distributed training —
dp_size: null # data‐parallel size
fs_size: null # FSDP size

# — Optimizer & Hyperparameters —
optimizer: dion2
rank_fraction: 0.5
scalar_opt: lion
mu: 0.95
weight_decay: 0.01
ortho_fraction: 0.5
adjust_lr: spectral_norm
lr: 0.02
21 changes: 9 additions & 12 deletions configs/dion_160m.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,33 +14,30 @@ num_iterations: 3000
warmup_ratio: 0.0
warmdown_ratio: 0.2

# — Optimizer & Hyperparameters —
rank_fraction: 0.125
mu: 0.95
weight_decay: 0.01
oversample: 1.25

# — Validation & Checkpointing —
val_loss_every: 125
val_tokens: 10485760
save_every: 0
val_tokens: 10485760

# — Weights & Biases logging —
wandb_project_name: gpt-train
wandb_job_name: null
no_wandb: false

# — Miscellaneous flags —
debug: false
no_compile: false

# — Distributed training —
dp_size: null # data‐parallel size
fs_size: null # FSDP size
tp_size: null # tensor‐parallel size
replicate_mesh_grad_sync: true

# — Miscellaneous flags —
debug: false
no_compile: false

# — Optimizer & Hyperparameters —
optimizer: dion
scalar_opt: lion
mu: 0.95
weight_decay: 0.01
ortho_fraction: 0.5
lr: 0.02
mixed_precision: true
13 changes: 5 additions & 8 deletions configs/muon_160m.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,14 +14,9 @@ num_iterations: 3000
warmup_ratio: 0.0
warmdown_ratio: 0.2

# — Optimizer & Hyperparameters —
mu: 0.95
weight_decay: 0.01

# — Validation & Checkpointing —
val_loss_every: 125
val_tokens: 10485760
save_every: 0
val_tokens: 10485760

# — Weights & Biases logging —
wandb_project_name: gpt-train
Expand All @@ -30,15 +25,17 @@ no_wandb: false

# — Distributed training —
dp_size: null # data‐parallel size
fs_size: null # FSDP size
tp_size: null # DO NOT USE TP for Muon
fs_size: null # FSDP size

# — Miscellaneous flags —
debug: false
no_compile: false
no_triton: false

# — Optimizer & Hyperparameters —
optimizer: muon
scalar_opt: lion
adjust_lr: spectral_norm
mu: 0.95
weight_decay: 0.01
lr: 0.02
13 changes: 5 additions & 8 deletions configs/normuon_160m.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,14 +14,9 @@ num_iterations: 3000
warmup_ratio: 0.0
warmdown_ratio: 0.2

# — Optimizer & Hyperparameters —
mu: 0.95
weight_decay: 0.01

# — Validation & Checkpointing —
val_loss_every: 125
val_tokens: 10485760
save_every: 0
val_tokens: 10485760

# — Weights & Biases logging —
wandb_project_name: gpt-train
Expand All @@ -30,15 +25,17 @@ no_wandb: false

# — Distributed training —
dp_size: null # data‐parallel size
fs_size: null # FSDP size
tp_size: null # DO NOT USE TP for NorMuon
fs_size: null # FSDP size

# — Miscellaneous flags —
debug: false
no_compile: false
no_triton: false

# — Optimizer & Hyperparameters —
optimizer: normuon
scalar_opt: lion
adjust_lr: spectral_norm
lr: 0.02
mu: 0.95
weight_decay: 0.01
Loading