Fix the distribution logic for the JAX trainer. #832
Conversation
1. eval() should also do a build for uninitialized metrics. 2. model.metrics should take care of uncompiled model.
keras_core/trainers/trainer.py
Outdated
| metrics.extend(self._metrics[:]) | ||
| if self._compile_metrics is not None: | ||
| if getattr(self, "_compile_metrics", None) is not None: |
There was a problem hiding this comment.
Why this change?
- The loss tracker should already be part of
metrics metricsshould definitely be called after__init__so_compile_metricsshould exist (if not that seems like a bug)
There was a problem hiding this comment.
I was hitting some test error for this. The loss_tracker is only inited after model.compile(), same for _compile_metrics .
This was causing issue when I try to access model.metric_variables for a uncompiled model, eg during model.predict().
There was a problem hiding this comment.
There was a problem hiding this comment.
I think you should be able to fix it by adding a if self.compiled: clause in _record_training_state_sharding_spec/etc to only look at optimizer variables and metric variables if the model is compiled, and not before.
There was a problem hiding this comment.
Oh, I think in this case, the if self.compiled should be added here, since _loss_tracker and _compile_metrics will only be populated after model.compile. Also the fix has to happen here, otherwise for a unacompiled model/layer, the model.metrics will raise an error which is a bad UX.
|
As far as I can tell the test failure isn't related. I'll look into it. |
The key fix here is too unwrap the jax.array from the trainable_variable/optimizer_variable before feed them to jax.jit train/eval function. Somehow JAX wasn't able to infer the sharding spec from the kerasVariable, which cause it to run with fully replicated weights.
Lift the data distribution logic out and put it outside of the jitted function. I noticed this issue when I debug the sharding setting.
Enforce the sharding constrain for the training state before the end of the jitted function. This will make sure the training state output have the same sharding step as the input. Otherwise it will cause jax to recompile the training function due to the sharding change.
I don't have a way to unit test this, since we can't easily step into a jitted function. The early test result from OPT2 model shows a great reduction of per-device memory usage, which means the distribution is actually working now.