Multiple iterations per training loop step?

[kovasb]

Hi,

In TF TPU usage a common pattern is to train for multiple steps per train call, https://www.tensorflow.org/guide/tpu#improving_performance_by_multiple_steps_within_tffunction

I was wondering how we can achieve a similar patten with objax. Do we need a custom version of the Jit class?

Along these lines it might be helpful to expand the docs on how these types of things work. Looking at the code, there seems to be some dance with temporarily replacing the module variables with their traced counterparts, and using side effects rather than return values to convey computation results.

Thanks!

[AlexeyKurakin]

Hi,

For training on TPU, I would suggest you to look at https://github.com/google/objax/tree/master/examples/classify/img/imagenet which was initially written for TPU, but also works on multi-GPU system. In particular, take a look at train_op method and self.train_op_parallel attribute of Experiment class.

For TPU or multi-GPU system, one have to use objax.Parallel instead of Jit to parallelize and compile ops. Jit will also work, but Jit will run all ops on single GPU or single TPU core.

Also as you correctly noticed, Objax uses trick with temporary variables replacing module variables to make OOP paradigm work with functional core of JAX. This works completely transparent on any computational device and there is no need to write any other custom version of Jit.

Regarding running multiple training steps per single train call. The above mentioned example does not do it, nevertheless I found it's performance just a little slower compared to my other TF2 implementation which performs several steps per train call. Note that I won't be providing specific performance numbers at right now, but we may consider doing more thorough benchmarking later. Also we may add example which performs multiple training steps per train call later as well.

[AlexeyKurakin]

So basically, if you're concerned about performance on TPU I would suggest to try existing examples with objax.Parallel. Likely performance will be satisfactory.

If you still need multiple training steps per single call of train op, we can consider adding such example later.

If you willing to contribute, feel free to add such example yourself. I suspect that adding regular for loop inside train_op may do the trick (i.e. it will be compiled to XLA while loop). However you need to figure out how to actually stream multiple batches of data into train_op

[AlexeyKurakin]

Also, I briefly looked at other JAX libraries (Haiku and Flax). Seems like all of them alway perform single training step per train_op. I suspect that either multiple training steps are unnecessary (i.e. don't provide any performance gains) or there are some issues with feeding multiple batches of data into single call of train_op compared to Tensorflow.

[AlexeyKurakin]

I found example of how to perform multiple training steps per one call to train function: https://github.com/mlperf/training_results_v0.7/blob/3dbb53064a6b79354c68a6832414b6536fee1a75/Google/benchmarks/resnet/implementations/resnet-research-JAX-tpu-v3-8192/train.py#L603

It seems like it require to manually maintain infeed queue and manually write on-device while_loop which all looks cumbersome. Given complexity of this setup, it makes sense to run benchmarks first before commiting to see how much performance we can gain from it.

[AlexeyKurakin]

@kovasb is my explanation provided any clarity of the question? Also would you be willing to do some prototyping / benchmarking of this change to see how useful it is for performance?

[AlexeyKurakin]

I'll close this issue for now. If we figure out that there is a need to have multiple iterations per training loop, then I'll reopen it.