Added training script for Gemma model

[vfdev-5]

What does this PR do?

• Added training script for gemma example based on lm1b_nnx example:
  • Added training code from l1mb_nnx
  • Added support distributed training via sharding, tested on 2 GPUs and TPU VM with 4 devices.
  • Added mixed precision config
  • Add support for multiple samples per sequence as in lm1b (nothing to modify in the attention layer, just use appropriate attention mask and provide shifted data)

Addresses #4740

• 2 GPUs training logs, gemma3-1b model config: link (1000 iters)
• TPU v4-8 training logs, gemma3-1b config: link (40000 iters)