[dont merge]. Diff 0.2 -> 0.3 - #2

Open

rakkit wants to merge 407 commits into

TrustLLMeu:trustllm-0.2from

rakkit:trustllm-0.3

This pull request is big! We're only showing the most recent 250 commits

Commits on Nov 17, 2025

Add export_dtype parameter to `convert_to_hf` function (#2041)
idoh
authored
22e959a

Commits on Nov 18, 2025

[compiler toolkit] Port joint_ac_pass from simplefsdp (#2051)
yiming0416
authored
3819737
Allow more arbitrary dataset selection

janEbert
authored and
rakkit
committed
f8ab834
Allow specifying arbitrary data files

janEbert
authored and
rakkit
committed
339d8dd
Add gradient accumulation test

janEbert
authored and
rakkit
committed
076d706
fix BitNet Llama's TP plan

janEbert
authored and
rakkit
committed
2652308
Add ByteTokenizer

janEbert
authored and
rakkit
committed
1fd11c9
remove llama-2 spec

janEbert
authored and
rakkit
committed
1af0dd0
Add byte-tokenized Llama-3 train spec

janEbert
authored and
rakkit
committed
77b17f6
Allow padding vocab size By default, pad to multiple of 128

janEbert
authored and
rakkit
committed
c873c72
Allow setting vocab size from tokenizer

janEbert
authored and
rakkit
committed
9002002
Add helper function for forward/param dtypes

janEbert
authored and
rakkit
committed
c812570
Expose additional data loader arguments

janEbert
authored and
rakkit
committed
bbefc8f
Add Byte tokenizer spec

janEbert
authored and
rakkit
committed
b7e854a
add bitnet opts, but disable it on job_config

janEbert
authored and
rakkit
committed
d2eff6b
Try to make BitNet work by patching PyTorch

janEbert
authored and
rakkit
committed
c0959b1
Refactor BitNet model code

janEbert
authored and
rakkit
committed
4c0d83a
Add QK normalization

janEbert
authored and
rakkit
committed
dd1b6cb
Add model config with QK normalization

janEbert
authored and
rakkit
committed
76dab56
Allow configuring norm layers

janEbert
authored and
rakkit
committed
c52d45d
Move BitNet model to separate file

janEbert
authored and
rakkit
committed
79bb06d
Add FineWeb data preset
rakkit
committed
48d59c5
Allow logging additional data
rakkit
committed
413004e
Add MoE model
rakkit
committed
29462d9
Add Scion optimizer
rakkit
committed
6833a20
Remove Scion arguments' defaults

ofivite
authored and
rakkit
committed
37f20a1
Add multi-token prediction objective

janEbert
authored and
rakkit
committed
81fde71
[mtp] Convert inputs and outputs to dictionaries

janEbert
authored and
rakkit
committed
afcb80b
Make MTP loss work with other loss functions

janEbert
authored and
rakkit
committed
56a4cad
Fix MoE loss with pipelining

janEbert
authored and
rakkit
committed
fb6624b
Refactor out sample packing

janEbert
authored and
rakkit
committed
28854b4
Add data mixing

janEbert
authored and
rakkit
committed
64715c2
Adapt arguments for dataset mixing

janEbert
authored and
rakkit
committed
f70c77e
Allow placing norms after every layer
rakkit
committed
20a1492
Support `norm_everywhere=True` model parallelism
rakkit
committed
eae92d8
Add activation offloading
rakkit
committed
42a8984
Allow separate seed for data

janEbert
authored and
rakkit
committed
0c4403f
Add windowed data shuffling

janEbert
authored and
rakkit
committed
2b2e197
Seed RNG unconditionally

janEbert
authored and
rakkit
committed
7b2ae33
Save job config upon Trainer creation

janEbert
authored and
rakkit
committed
2bd927e
Save model args

janEbert
authored and
rakkit
committed
2e621f2
Use more legible JSON format
rakkit
committed
d3f924f
Also dump training config as TOML
rakkit
committed
d8eb36e
Add KV cache

janEbert
authored and
rakkit
committed
2bec9a9
Expose more WandB settings

ofivite
authored and
rakkit
committed
73cc61e
Allow configuring separate param groups

ofivite
authored and
rakkit
committed
5782c09
Explicitly pass `device_id` to `torch.distributed`

janEbert
authored and
rakkit
committed
198a31a
Allow logging parameter and update norms

ofivite
authored and
rakkit
committed
5ca2be4
Fix MoE norm logging
rakkit
committed
3928627
Allow logging learning rates per parameter group

ofivite
authored and
rakkit
committed
519b406
Add config for norm experiments

ofivite
authored and
rakkit
committed
0cf1558
Allow model initialization configuration

janEbert
authored and
rakkit
committed
44041d7
Refactor updating model config from job config

janEbert
authored and
rakkit
committed
65925e5
Expose initialization configurations

janEbert
authored and
rakkit
committed
40b905d
Use muP-like initialization in norm experiments

janEbert
authored and
rakkit
committed
4b0e244
Implement muP hyperparameter width multipliers

ofivite
authored and
rakkit
committed
d9a7acd
Disable weight decay for norm experiments

janEbert
authored and
rakkit
committed
409cdcc
Support Scion initialization scheme

ofivite
authored and
rakkit
committed
78751e7
Add debug model configs for scaling tests

ofivite
authored and
rakkit
committed
6adaf01
Change norm experiment momentum to Scion's default

ofivite
authored and
rakkit
committed
85bc8ed
Change model scaling debug configs

ofivite
authored and
rakkit
committed
8b658ca
Prefer non-light Scion for norm experiments

ofivite
authored and
rakkit
committed
be7e14d
Allow disabling gradient clipping

ofivite
authored and
rakkit
committed
edfb59d
Disable gradient clipping for norm experiments

ofivite
authored and
rakkit
committed
2837df5
Fix saving/loading for light Scion

janEbert
authored and
rakkit
committed
732d330
Fix expert parallelism
rakkit
committed
5f3b30c
Expose `start_pos` argument

janEbert
authored and
rakkit
committed
ddba03b
Add distributed Scion
rakkit
committed
3f595e9
Fix generator argument handling

janEbert
authored and
rakkit
committed
11e6c16
Pass `std` argument to orthogonal init

janEbert
authored and
rakkit
committed
a3dff1a
Fix scaling debug model configs

ofivite
authored and
rakkit
committed
2fb7bc1
Allow logging condition numbers

ofivite
authored and
rakkit
committed
2d46c20
Do not include epsilon in sqrt for normalization

ofivite
authored and
rakkit
committed
e75d2d6
Fix light Scion and refactor

ofivite
authored and
rakkit
committed
c7d2842
Add scaled orthogonal initialization

ofivite
authored and
rakkit
committed
ce26df2
Fix Scion initialization

ofivite
authored and
rakkit
committed
d55db43
Replace spectral with RMS->RMS norm logging

ofivite
authored and
rakkit
committed
fefbecd
Improve condition number numerical stability

ofivite
authored and
rakkit
committed
b545956
Make scaling-relevant configs part of model flavor

ofivite
authored and
rakkit
committed
2e9c6ea
Allow disabling depth-wise residual init scaling

ofivite
authored and
rakkit
committed
8f97e72
Do not include epsilon in sqrt for normalization

janEbert
authored and
rakkit
committed
690fa16
Fix distributed Scion
rakkit
committed
7229cdf
Align APIs
rakkit
committed
d7ea4fe
Add padding token support to byte tokenizer

janEbert
authored and
rakkit
committed
8536216
Add padding token support to model

janEbert
authored and
rakkit
committed
c7290b8
Take `eos_id` from tokenizer if available

janEbert
authored and
rakkit
committed
499069b
Improve optimizer logging
rakkit
committed
969ee72
Allow configuring Scion backend
rakkit
committed
7ceee8c
Expose Scion backend config
rakkit
committed
d23279b
Force SVD dtype to float32
rakkit
committed
587ee12
Adapt MoE debug config to Llama-3 and Scion
rakkit
committed
b4a4617
Support MoE router customization
rakkit
committed
7df9af0
Add image domain orthogonal initialization
rakkit
committed
b0eda73
Add image domain norm for Scion
rakkit
committed
5ea5e57
Fix style
rakkit
committed
fa35fd9
Add debug model with QK norm
rakkit
committed
7930f92
Do not modify routing parameter group for Scion
rakkit
committed
769d165
Disable experimental activation checkpointing
rakkit
committed
266a966
Update Scion implementation capabilities
rakkit
committed
c1ce271
Fix distributed Scion query for light version
rakkit
committed
f6a4f8b
Allow configuring LR when resuming from checkpoint
rakkit
committed
0bc6216
Allow configuring whether LRs are loaded

janEbert
authored and
rakkit
committed
bb54204
Fix distributed Scion param group configuration
rakkit
committed
197fb0a
Fix line lengths

janEbert
authored and
rakkit
committed
3fb2a59
Add baseline configs
rakkit
committed
6104063
Fix target shape indexing
rakkit
committed
c860620
Refactor Scion backend selection
rakkit
committed
8827d36
Do not add empty parameter groups
rakkit
committed
ba4fc11
Fix RNG state loading
rakkit
committed
d29c99b
Hackily force process exits upon completion
rakkit
committed
d6eae15
Hackily fix circular import error
rakkit
committed
6f8d194
Refactor optional Attention norms
rakkit
committed
7270b9b
Change Scion parametrization settings
rakkit
committed
dfef8f5
Fix parameter norm calculation for MoE
rakkit
committed
ad12138
Fix Scion with MoE
rakkit
committed
4bfaf66
Fix bug
rakkit
committed
bd804fc
Skip empty local tensors
rakkit
committed
8900c4b
Fix missing dependency for activation offloading
rakkit
committed
758b0ee
Fix typing and adjust argument names

janEbert
authored and
rakkit
committed
495338d
Move `_wrap_orthogonal` into global scope

janEbert
authored and
rakkit
committed
bf527eb
Refactor function selection in `build_init_fn`
rakkit
committed
7e4fb7d
Add no-communication version of `orthogonal_`
rakkit
committed
7f97d10
Handle `None` gradients
rakkit
committed
7a44aef
Refactor module name clean-up function

janEbert
authored and
rakkit
committed
3688cc6
Refactor norm functions

janEbert
authored and
rakkit
committed
4ee293d
Refactor Scion
rakkit
committed
af282ec
Refactor parameter norm calculations
rakkit
committed
eb05c33
Refactor orthogonal init
rakkit
committed
9c69e53
Add Polar Express backend
rakkit
committed
852b7f3
Add additional norms
rakkit
committed
054feaf
Add SSNorm
rakkit
committed
8ae25e8
Fix SSNorm treatment with Scion
rakkit
committed
811c63e
Allow setting `depth_init` from job config
rakkit
committed
2dc586d
Warn about regexp mismatch more prominently
rakkit
committed
d6ed876
Add Scion gradient debugging mode
rakkit
committed
ddd57e7
Allow configuring model residual scaling
rakkit
committed
0358ed0
Allow selection of logged norms
rakkit
committed
bce1df8
Skip model weight initialization if possible

janEbert
authored and
rakkit
committed
177ce3e
Fix and refactor distributed Scion
rakkit
committed
c33316d
Log expert metrics if available
rakkit
committed
ceda66a
Add more model flavors
rakkit
committed
44df7ef
Add debug model flavors
rakkit
committed
b2b3317
Refactor MoE optimizer modifications
rakkit
committed
c0a366f
Improve readability

janEbert
authored and
rakkit
committed
0101c62
Fix MoE model args
rakkit
committed
62aa68b
Refactor MoE pipelining
rakkit
committed
5e190f2
Refactor setting layer ID
rakkit
committed
abb3468
Move class to instance variables

janEbert
authored and
rakkit
committed
54e9d17
Fix MoE weight initialization
rakkit
committed
8af93ef
Refactor expert parallelism
rakkit
committed
18d5b50
Remove temporary patches
rakkit
committed
8973665
Adjust MoE parallelization to upstream Llama-3
rakkit
committed
f693ffe
Refactor/fix MoE
rakkit
committed
8a722b5
Log additional througput-related metrics
rakkit
committed
17818f3
Remove legacy code
rakkit
committed
7f8a244
Delete unused code
rakkit
committed
422869c
Allow configuring activation function
rakkit
committed
7149767
Add more MoE job configurations
rakkit
committed
e7a45bb
Allow configuring MoE activation function
rakkit
committed
48e5e68
Remove redundant pre-defined models
rakkit
committed
89a9491
Remove initializing all MoE experts the same
rakkit
committed
8ec93de
Fix passing MoE arguments to sublayer
rakkit
committed
6cbc5a5
Remove option to use bias for routing
rakkit
committed
1526633
Make sequence-wise auxiliary loss staticmethod
rakkit
committed
2429f61
Fix docstring placement

janEbert
authored and
rakkit
committed
b6c291f
Add zero-mean MoE router bias update options
rakkit
committed
c35d72d
Fix `norm_everywhere=True` with tensor parallelism
rakkit
committed
6f5b02e
Apply PP compatibility hacks
rakkit
committed
08f96a9
Optimize MoE metrics logging
rakkit
committed
c936185
Make Scion work with biases and weight decay
rakkit
committed
64c797c
Add optimized zeropower backend Triton kernels
rakkit
committed
593ba0d
Disable MoE aux loss by default
rakkit
committed
80dd38f
Fix Polar Express
rakkit
committed
2e2887d
Fix MoE bias update scale factor
rakkit
committed
ec1dd2d
Log MoE router data only on DP+CP rank 0

janEbert
authored and
rakkit
committed
7ae36b9
Allow more selective logging

janEbert
authored and
rakkit
committed
db89486
Allow enabling logging for all first PP stages
rakkit
committed
425d92f
Log Scion norms only on DP+CP rank 0
rakkit
committed
5a2415c
Do not pad vocab size by default

janEbert
authored and
rakkit
committed
4f47c03
Allow setting vocab size via CLI

janEbert
authored and
rakkit
committed
eb59b4d
Fix MoE logging
rakkit
committed
20b9e72
Transpose grouped expert weights
rakkit
committed
9cef2c6
Prefer floor division for integer metrics
rakkit
committed
8d4f82b
Add forced uniform MoE load balancer prototype
rakkit
committed
27dfd5f
Rename "feed_forward" to "moe" when using MoE
rakkit
committed
6c16352
Sync MoE parallelization upstream changes
rakkit
committed
2d4490e
Handle unsupported MoE+Scion setting
rakkit
committed
7dc850f
Fix `depth_init` string comparisons

janEbert
authored and
rakkit
committed
3ba66a2
sync some upstream changes. and rename .ffn of moe to .moe align up steram, re-use the apply_FSDP from llama4 for our model
rakkit
committed
84f4099
add conv code for scion, correct name to rms->inf
rakkit
committed
0d91a00
Allow more arbitrary dataset selection

janEbert
authored and
rakkit
committed
eecf257
Allow specifying arbitrary data files

janEbert
authored and
rakkit
committed
687ce31
Add multi-token prediction objective

janEbert
authored and
rakkit
committed
7091669
Refactor out sample packing

janEbert
authored and
rakkit
committed
ed6bd9a
refactor scion init, disable scion-light until we check it
rakkit
committed
ebb988d
dont register optim hook for none moe model_parts, hack for PP+MoE
rakkit
committed
06d839c
fix some bugs
rakkit
committed
4cd7f54
improve readability of depth init
rakkit
committed
f46f412
improve readability of depth init
rakkit
committed
51756a7
rename it to disco.py
rakkit
committed
54ed8f1

Commits on Nov 19, 2025

from upstream torchtitan/datasets/hf_datasets.py -> torchtitan/hf_datasets/text_datasets
rakkit
committed
81f8e55
fix pyproject for pre-commit
rakkit
committed
9b5934f
fix train.py and args after rebase. and refactor args
rakkit
committed
2ff168a
disable compile in lmo's normlize
rakkit
committed
bdca787
formating
rakkit
committed
9043220
formating
rakkit
committed
5dfef2b
refactor llama3 about args & config
rakkit
committed
e46db6e
refactor moe about args and config. taking upstream changes on moe
rakkit
committed
e671ca5
make sure PP not fail when splites accidently set to [ ]
rakkit
committed
0321e3b
fix a strange import error and train spec cannot find modules due to old register api
rakkit
committed
1e54d4b
fix a strange import error and train spec cannot find modules due to old register api
rakkit
committed
9944c38
decouple moe's auxloss to make main loss be only CE, and rename aux loss to load_balance_loss
rakkit
committed
416bb86

Commits on Nov 20, 2025

Commits on Nov 26, 2025

disable rope scaling for llama3 by default
rakkit
committed
9b120e5
record models' dim-in and dim-out
rakkit
committed
67b846f
save HF's format of config when convrt DCP to HF ckpt
rakkit
committed
fecaa44
add MoEllama's DCP <-> HF convert
rakkit
committed
b118d11
add HF assests for llama and MoEllama, when convert DCP to HF automatically copy these code to HF ckpt
rakkit
committed
d937084
set checkpoint.last_save_model_only to *False* by default, to avoid we lost all the checkpoint
rakkit
committed
15cea0f

Commits on Nov 27, 2025

add data mixing scheduler
rakkit
committed
9e3da55
fix a bug that attention fail for 2x query head
rakkit
committed
d0849fa
formating
rakkit
committed
a3eb33e

Commits on Nov 28, 2025

improving logging params and flops
rakkit
committed
0c3abd8
make prepare momentum works faster
rakkit
committed
55f247d

Commits on Dec 1, 2025

add proxy model
rakkit
committed
1b9a56c
improve logging
rakkit
committed
647d3e8

Commits on Dec 5, 2025

fix entropy logging
rakkit
committed
ce9d75e

Commits on Dec 6, 2025

fix moe's ddp parallel func
rakkit
committed
5c670a4
disable skip weights init to makesure non persistent can be init from ckpt
rakkit
committed
8ad768e

Commits on Dec 9, 2025

improve HF convert
rakkit
committed
cdffae5

Commits on Dec 10, 2025

fix load balance aux loss
rakkit
committed
0c6a480

Commits on Dec 11, 2025

fix scion norm with compile, see https://github.com/SDLAML/disco/pull/1
rakkit
committed
6531a40
refactor dataset a bit to clean legacy code
rakkit
committed
add6147

Commits on Dec 13, 2025

add sft dataset and patch for plelas synt data
rakkit
committed
d56deb0

Commits on Dec 15, 2025

fix hf adapter for moe
rakkit
committed
8519e32

Commits on Dec 17, 2025

better data mix logger
rakkit
committed
2cef47a
fix logger in dataset
rakkit
committed
156c8c7

Commits on Dec 22, 2025

improve and fix data sample
rakkit
committed
436443e
add few more logging
rakkit
committed
c2d3b10

Commits on Dec 24, 2025

cleanup sft dataset, will move to experiment later
rakkit
committed
839f8f2
revert compile in disco
rakkit
committed
b320fa0
allow to set pre-fetch
rakkit
committed
dccec36
improve and fix seed for data mix
rakkit
committed
e418299

Commits on Dec 25, 2025

fix sum_data_sampled all-reduce
rakkit
committed
305a020

Commits on Dec 31, 2025

fix and improve logic of exhuasted datasets
rakkit
committed
bedf495
fix dataset state and log out the actual sample ratio
rakkit
committed
6c09335

Commits on Jan 3, 2026

fix and improve datalaoder
rakkit
committed
4290c95

Commits on Jan 6, 2026

fix on num_sampled_per_dataset state in loader
rakkit
committed
64cc998
improve hf model convert
rakkit
committed
de0ce0e