feat(training): multi-worker training

[clement-bonnet]

• adapt training code to run on multiple workers (e.g. TPU v4-16 and more)
• [x] tested on CPU
• [x] tested on TPU v4-8
• [x] tested on TPU v4-16 (2 workers)

[clement-bonnet]

CPU

TPU v4-8

TPU v4-16

TPU v4-16 is indeed twice as fast ("steps per second" metric) as TPU v4-8.