Add an example of training a tabular model on multiple GPUs

[akihironitta]

Adds an example for training a model on multiple GPUs.

Usage

$ python examples/trompt_multi_gpu.py

Changes

To enable it, this PR modifies Trompt accordingly:

- out = Trompt(...).forward_stacked(tf)  # forward_stacked is removed
+ out = Trompt(...)(tf)
  assert out.size() == (batch_size, num_layers, num_channels)

Some highlights in the script

• Unlike examples/trompt.py, it computes training accuracy batch-wise instead of computing it at the end of each epoch to save time (although model parameters change over steps within each epoch).
• It reduces losses and metrics with all_reduce and torchmetrics's API across all ranks at the end of each epoch.
• It avoids unnecessary device synchronisations within each epoch, e.g., by not calling float(cuda_tensor), and by setting repeat_interleave(..., output_size=...).
• ... (I'm happy to elaborate if anything in the script is unclear.)

Benchmark results from --dataset jannis on g6.12xlarge with four L4 GPUs

	four GPUs	single GPU
time per training step	0.725 seconds	0.741 seconds
time per training epoch	29 seconds	117 seconds
test accuracy	80.23 %	80.29 %

[weihua916]

Let's make sure this change does not break the example code.

[akihironitta]

Confirmed the change doesn't break these scripts across all supported task types:

examples/trompt.py
benchmark/data_frame_benchmark.py
benchmark/data_frame_text_benchmark.py