On Heterogenous Cluster

[hezq06]

Our lab has 3 servers (bought at different years) with a speed ratio 10:9:5. To maintain load balance, I'm thinking of tweaking the batch size also to be 10:9:5. This seems to be a simple idea (being able to set batch size for each node separately and use that number as a weight during the parameter update stage). However, I cannot find an easy solution for implementing that with DeepSpeed. Does anyone know a way?

[jeffra]

Hi @hezq06, thanks for your question this is interesting. I have never tried this, but I believe you should be able to do this mostly outside of DeepSpeed between your dataloader and making sure you get the 10:9:5 split you desire for a batch before passing it to model(batch). DeepSpeed will take over at the forward/backward/step phases of training but the batch size prep is done before forward. I would then suggest fixing the train_micro_batch_size_per_gpu and gradient_accumulation_steps values in your ds_config to a single value. The micro batch size will only be used for calculating SamplesPerSec in this case, which won't really be accurate for you but otherwise i think DeepSpeed shouldn't care about this heterogenous setup.

One thing to note, I would make sure to align torch versions across your cluster and ideally cuda versions as well.

[hezq06]

Hi @jeffra, thank you for your reply. Handling heterogeneous clusters outside DeepSpeed with dataloader taking care of the 10:9:5 split sounds really reasonable to me. However, I've got another concern. Since I have no idea about how deepspeed gathers gradient information from all nodes and calculates parameter update, one possibility is that all batched gradient information is first collected and then averaged, and another possibility is that gradient information is first averaged on each node, then gathered, and then averaged again. Possibility one is quite unlikely since the communication overhead is huge. Possibility two seems more likely, however, the speed-10-machine-generated gradient would then be treated equally with the speed 5 machine, while in reality, the speed 10 machine would have a gradient that is less noisy. As a result, the calculated gradient is sub-optimum. It would be better also to communicate the batch size but I'm not sure if deepspeed already handles this.

P.S. Thank you for your extra note. Actually, I once begin with a different torch and cuda version, but nccl just can't successfully initialize. After aligning the versions, things work well.