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.

fix: style

feat: add more config of distributed_shampoo

feat(train): refactor learning rate params

fix(train): handle seed_dataset

feat: refactor TrainingArguments

fix: push_to_hub deprecated

fix(inference): use float32 + flatten logits

feat: update inference pipeline

style: isort

style: apply to distributed_shampoo

feat: add best_effort_memory_usage_reduction

feat: update distributed_shampoo

fix: weight decay Adam + speed logging

feat: add micro config

fix: shampoo -> distributed shampoo

feat: update params

feat: add shampoo optimizer

feat: update sweep

feat: create config files

feat: allow abstract_init

fix: typo

feat: log more metrics

feat: shard by host is optional

feat: load data first

feat: display local TPU's