feat: vmap optimizer (#166)

feat: scan layers + gradient checkpointing (#161)

feat: better multi-node support (#158)

feat(model): allow bias (#152)

feat(train): google-cloud-storage is optional

feat(train): rename logged config

feat: implement transformer variants (#144)

feat(train): log norm and histograms (#143)

feat(data): super conditioning (#141)

feat: support pod (#139)

fix: no gradient checkpointing for new model

feat: no gradient checkpointing for params init

fix(train): consider schedule offset

feat(train): local jax cache

feat: add bucket reference to artifact

style: lint

feat: handle gradient checkpointing

feat: load from bucket

feat(train): save to bucket

feat: reduce artifact space + offset step

feat: restore weights on CPU

feat(train): simplify tokenizer loading

feat(train): use compilation cache

feat: log num_parameters early

feat(train) - handle multiple nodes (#130)

feat: handle model parallel

feat(train): more custom x-axis

fix(train): opt_state_shape for distributed_shampoo

feat(train): split artifact into model/state

feat(train): another 25% faster

feat(train): overhead from 70% to 1% 🥳

feat(pjit): follow t5x style

fix(train): grads spec

feat(train): improve pjit speed

fix(train): consider correct batch size

feat(train): custom start_preconditioning_step

feat(train): handle distributed_shampoo in pjit

feat(train): distributed_shampoo with pjit

fix style

feat(train): restore opt_state efficiently

feat(model): clean way to load on cpu

feat(train): load model on CPU

feat(train): different rng per node

feat(train): no batch dimension with pjit

feat(train): progress on pjit

feat(train): start pjit support

Load from wandb artifact (#121)

Use DalleBartTokenizer. State restoration reverted to previous method:

fix(train): variable not defined

feat(train): cleanup args