deepspeedai / deepspeedai/DeepSpeed

Question about using Autotuner with ZeRO and tensor parallelism

Open
#6,796 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
43.1k
Forks
5k
Avg merge
4d 15h
Merged PRs (30d)
112

Description

I’m reading through the Autotuner code and found this function:

https://github.com/microsoft/DeepSpeed/blob/f743feca033515fdded50a98093da5a48eb41e74/deepspeed/autotuning/autotuner.py#L278-L302

It computes

total_gpus = self.exp_num_nodes * self.exp_num_gpus

based on the autotuning config. If ZeRO is enabled, then based on which stages are enabled, optimizer_mem, gradients_mem, and/or params_mem get sharded across the GPUs. But then if self.mp_size() (for tensor parallelism, right?) is greater than 1, then the total memory usage is divided again by the amount of tensor parallelism. So if ZeRO and tensor parallelism are both enabled, this is double-dipping, right? With N GPUs, we can’t get the per-GPU memory usage any smaller than 1/N. I’m not sure if

  1. there’s a bug here
  2. if the value of the num_gpus flag supposed to be reduced by the amount of tensor parallelism, or
  3. if I’m not understanding this correctly.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in deepspeed/autotuning/autotuner.py at lines 278-302 and trace the ZeRO memory calculations alongside self.mp_size() for tensor parallelism. Compare the resulting per-GPU memory estimate when both features are enabled; done means determining whether memory is divided twice and identifying the intended handling of num_gpus.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.