… separately from trainer

… logging (#367)""

This reverts commit 137bd4b.

* Tweaks to improve behavior on slurm clusters (#373)

* giving docker a shot

* wip prepare_cache

* fix deps

* push slurm stuff into logical_cpu_count

* add release for build so it's the default

* roll slurm nodelist expansion ourselves

* don't use cluster stuff for initializing ray

* cleanup

* cleanu

* no levanter slurm cluster

* fix validation set when it's an HF dataset

* do basic logger initialization up front

* add more logging for distributed init

* pretty sure this was breaking multihost levanter

* better slurm nodename

* use the LevanterSlurmCluster thing again

* fix commented out line

* less obnoxious logging

* wip

* initial commit of hero

* switch to autodiff hvp

* make test run anywhere

* remove stub adam thing

* fix missing returns in text.pyu

i don't understand how this code made it in/wasn't caught

* fix "replicated" arrays in named_pjit

* support hessian updates for hero in gpt2_example

* track time for hessian computation

* attempt to add some logging from inside jax

* stupidly was squaring the LR

* make global batch loader invariant to reorderings of how jax.array_from_callback call the callback (#132)

* default parallelism

* add decay schedule for gamma

* rename hero to sofia

* wip

* mucking around with sharding b/c of some weird behaviro on cpu

* let's see if this works ons TPU

* fix sofia

* fix test for named_jit

* wip

* wip making sophia configurable

* wip

* wip

* mdakd

* wip

* fix arg name

* update configs

* explicit sophia_g config

* wip

* not sure how this didn't get committed

* fix sophia-g (dumb dtype issue)

* try logging to see what the hold up is

* sigh

* try to jit the eval loader and see if that fixes it

* refactor loaders to have more in common, should hopefully improve perf

* fix sophia test

* update sophia large configs

* missed a sofia

* tweak lr

* small_fast sophia_h config

* got sophia_h test working again

* please pre-commit

* fix alpaca ocnfig for llama

* almost there for sophia-h

* switch to 0.4.19 b/c there's a bug in jaxlib that sophia triggers

* ok, see if this is correct-ish

* try a fancier config

* put clipping into the sophia scale so it's easier to manage

* try clipping the hessian on a per-component basis?

* more clip

* fix tests by default to adam

* try a full gpt2 small run

* add tag

* remove clip

* wip

* bump version

* disable logging for now

* try this?

* cleanup, delete old optimizerconfig

* reorganize optim package

* remove type: adam stuff since it's default

* forgot the hessian_update in the migration

* fix LR being reported as 0

* oops

* allow_val_change in wandb

* sanity check epochs

…traction