feat: reduce artifact space + offset step

feat(demo): update reference

feat: restore weights on CPU

feat(train): simplify tokenizer loading

feat: cleanup notebook

feat(demo): uncomment pip install

feat: improve inference demo

feat(train): use compilation cache

fix: style

Merge branch 'main' of https://github.com/borisdayma/dalle-mini into main

feat: log num_parameters early

fix: distributed shampoo class

feat: update distributed_shampoo

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: update distributed_shampoo + fix None spec

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)

feat(train): update sweep config

Use DalleBartTokenizer. State restoration reverted to previous method:

fix(train): variable not defined

feat(train): cleanup args

refactor(train): cleanup

feat: custom gradient accumulation

fix: style

feat(train): use MultiSteps for gradient accumulation

Accept changes suggested by linter.

Update help string for `model_name_or_path`.

Update `resume_from_checkpoint` to use `from_pretrained`.

Load tokenizer associated to the model checkpoint, if possible.

Use model configuration unless a specific one is supplied.