[FEATURE] TPU Training on a multiple nodes (bits_and_tpu branch)

[eliahuhorwitz]

Hey,
Is it possible, and if not, is it complicated to support training on a TPUv3 VM of 32 machines?

Thanks,
Eliahu

[rwightman]

@eliahuhorwitz it's on the todo list but I have not tried it, I assume it will just involve changing how the training launches but I'm not 100% sure.

[dedeswim]

Hi! I tried it and it works without issues. I just followed the instruction on XLA's README:

1. I created a container with a bunch of Python packages I need, and has this shape:

FROM gcr.io/tpu-pytorch/xla:r1.10_3.8_tpuvm

RUN pip install --upgrade pip

RUN mkdir /<your-dir>/
WORKDIR /<your-dir>/

# I have `git+https://github.com/rwightman/pytorch-image-models.git@fafece230b8c8325fd6144efbab25cbc6cf5ca5c`
# in my `requirements.txt`. This is a specific commit from `bits_and_tpu`, but I guess also `@bits_and_tpu` should work.
COPY requirements.txt requirements.txt
RUN pip install -r requirements.txt
RUN pip install wandb

# Assuming that you are in the directory with your code
COPY . .

2. I built this docker image and tagged it with some name of my choice. Then I pushed it to the Google Container Registry (the instructions are here.
3. After setting up docker on all the pods (more info here), I pulled this image from all the workers of my pod with gcloud alpha compute tpus tpu-vm ssh vits-robustness-pod --zone us-central2-b --worker=all --command "docker pull gcr.io/path-to-my-image"
4. I launched distributed training with

python3 -m torch_xla.distributed.xla_dist --tpu=${TPU_NAME} --restart-tpuvm-pod-server --docker-image=gcr.io/path-to-my-image --docker-run-flag=--rm=true --docker-run-flag=--shm-size=40GB --docker-run-flag=--mount=type=bind,source="$(pwd)"/output,target=/<your-dir>/output -- python launch_xla.py --num-devices 8 train.py <data-dir> --output /<your-dir>/output --use-mp-loader --model <your-model> --other-args

This worked on a v4 TPU (both 32 and 64 cores), and I guess will work also on a v3. In the directory of my code, I have the same launch_xla.py file as the one in the bits_and_tpu branch, and a custom version of the train.py in bits_and_tpu (I had to customize it because I have to do adversarial training). The --use-mp-loader flag is not in the original file, I have added as a bool argument, which, if True does the following at the end of the setup_data function:

if args.use_mp_loader and dev_env.type_xla:
        import torch_xla.distributed.parallel_loader as pl
        assert isinstance(loader_train, fetcher.Fetcher)
        assert isinstance(loader_eval, fetcher.Fetcher)
        loader_train.use_mp_loader = True
        loader_train._loader = pl.MpDeviceLoader(loader_train._loader, dev_env.device)
        loader_eval.use_mp_loader = True
        loader_eval._loader = pl.MpDeviceLoader(loader_eval._loader, dev_env.device)

I am not sure this is necessary. I did it just because in the distributed training examples on PyTorch XLA's repo they were using it