In the user training code, instead of directly specifying the process rank, there is a need to raise the level of abstraction by specifying the stage. For instance, instead of the rank of world_size - 1, it should be possible to refer to it as the last_stage.
In the user training code, instead of directly specifying the process rank, there is a need to raise the level of abstraction by specifying the stage. For instance, instead of the rank of world_size - 1, it should be possible to refer to it as the last_stage.