pytorch
Hpc setup
HPC type setup ( e.g. mpirun , mpiexec) currently requires code change in train.py ( or elsewhere) . Adding couple of optional additional args ( e.g. rank and local rank) to toml file along with a simple function in train.py removes the need to make those changes.
Aren’t those environment variables set automatically when using torchrun command?
As I mentioned above, this is for non-torchrun launchers like mpirun/mpiexec etc. Torchrun is very specific to torch, HPC environments run many other applications which can not use torchrun. HPC environments are optimized for whatever launcher they were designed for. In those environments, torchrun usually does not work or has sub-optimal performance.
@tianyu-l and @fegin , thanks for your replies. Let me try to address the questions one by one below.
- Setting variable in run_train.sh : It does not work because of a subtle reason. The job launchers like torchrun/mpirun/mpiexec set the local and global rank variable uniquely in the shell/process of each rank. Hence, setting rank variables outside the launcher will make them all same or non-unique.
- LOCAL_RANK=some_rank to launch command : same as 1 above .
- Extending-jobconfig: it is certainly possible to make it work, but it would be far more complex than what is needed. Also, it would require maintenance of a separate piece of code elsewhere. Different launcher's already uniquely set the ranks variables per-rank process/shell . All that is left to be done is copy them to the variables that TorchTitan is expecting.
@TJ-Solergibert , it has the same issue as 1 above, if I understand the script.
In torch distributed programs you have to set 5 environment variables:
MASTER_ADDR,MASTER_PORT&WORLD_SIZEwhich need to be COMMON for all the processes. This is done hereRANK&LOCAL_RANKwhich are PARTICULAR for every process. Recall that we have one process por GPU. This is done here.
In this example you can consider srun as mpirun and you set the number of processes through this 2 settings.
@TJ-Solergibert , thanks for pointing out the following line for SLURM launcher srun
srun $SRUN_ARGS bash -c "RANK=\$SLURM_PROCID LOCAL_RANK=\$SLURM_LOCALID $CMD" 2>&1 | tee -a $LOG_PATH .
The above also would benefit from simplification I am proposing - see below.
srun $SRUN_ARGS $CMD 2>&1 | tee -a $LOG_PATH
@githubsgi You can check https://github.com/pytorch/torchtitan/blob/main/docs/extension.md#extending-jobconfig. This should meet your goal instead of adding new ones to the main JobConfig.
@fegin , I have already addressed the complexity of that approach above. In fact, the approach suggested by @TJ-Solergibert above is simpler and more maintainable. The 2 environment variables I am proposing is applicable to any launcher and thus makes TorchTitan code very portable,