train.py

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
print(f"trainer got local_rank {args.local_rank}")
import torch
import torch.distributed as dist

torch.cuda.set_device(args.local_rank)
dist.init_process_group(backend="nccl", init_method="env://")
print(f"trainer {args.local_rank} initialized process group")

model = torch.nn.Linear(1, 1)
model = model.to(args.local_rank)
ddp = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank]
        )
print(f"trainer {args.local_rank} done initializing DDP")
import sys
print("stderr", file=sys.stderr)