deepspeedai / deepspeedai/DeepSpeed

no overlap while performing stage-2 zero dp training

Open
#4,599 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

Describe the bug
No overlap while performing stage-2 zero dp training according to torch profile:
image
I expect backward kernels should overlap with nccl kernels.

To Reproduce

we run it on 2 node , one P100(16GB) for each
ds config like:

  "zero_optimization": {
      "stage": args.stage,
      "allgather_partitions": True,
      "reduce_scatter": False,
      "allgather_bucket_size": 10000,
      "reduce_bucket_size": args.bucket_size,
      "overlap_comm": True,
      "contiguous_gradients": False,
      "cpu_offload": False
  },

notice overlap_comm=True

training script like below

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 1024)
        self.fc2 = nn.Linear(1024, 256)
        self.fc3 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        if args.moe:
            for layer in self.moe_layer_list:
                x, _, _ = layer(x)
            x = self.fc4(x)
        else:
            x = self.fc3(x)
        return x


net =  Net()

parameters = filter(lambda p: p.requires_grad, net.parameters())
model_engine, optimizer, trainloader, __ = deepspeed.initialize(
    args=args, model=net, model_parameters=parameters, config=ds_config)

local_device = get_accelerator().device_name(model_engine.local_rank)
local_rank = model_engine.local_rank

# For float32, target_dtype will be None so no datatype conversion needed
target_dtype = None
if model_engine.bfloat16_enabled():
    target_dtype=torch.bfloat16
elif model_engine.fp16_enabled():
    target_dtype=torch.half


import torch.optim as optim

criterion = nn.CrossEntropyLoss()
tbname = 'comm' if "comm_simulator" in ds_config else 'normal'
bs = args.batch_size


def run(profiler=None):
    for _ in range(10):  # loop over the dataset multiple times
        # fake data and labels simulating cifar dataset
        inputs, labels = torch.ones([bs,3,32,32]), torch.ones([bs, 10])
        if "comm_simulator" not in ds_config:
            inputs, labels = inputs.to(local_device), labels.to(local_device)

        if target_dtype != None:
            inputs = inputs.to(target_dtype)

        outputs = model_engine(inputs)
        loss = criterion(outputs, labels)
        model_engine.backward(loss)
        model_engine.step()

        if profiler:
            profiler.step()


with torch.profiler.profile(
    schedule=torch.profiler.schedule(
        wait=1, # During this phase profiler is not active.
        warmup=1, # During this phase profiler starts tracing, but the results are discarded.
        active=6, # During this phase profiler traces and records data.
        repeat=2), # Specifies an upper bound on the number of cycles.
    profile_memory=True,
    with_stack=True,
    on_trace_ready=torch.profiler.tensorboard_trace_handler(f'tb', worker_name=f'stage{args.stage}-bs {args.batch_size}-bucketsize {args.bucket_size}')
) as profiler:
    run(profiler)

print('Finished Training')

ds_report output

--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  please install triton==1.0.0 if you want to use sparse attention
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch version .................... 1.13.1
deepspeed info ................... 0.11.1, unknown, unknown
torch cuda version ............... 11.6
torch hip version ................ None
nvcc version ..................... 11.6
deepspeed wheel compiled w. ...... torch 1.13, cuda 11.6
shared memory (/dev/shm) size .... 93.74 GB

System info:

  • OS: Ubuntu 20.04.5 LTS
  • GPU count and types : two machines with P100 each
  • Interconnects (if applicable): two machines connected with 10 Gbps IB
  • Python version:3.8.18

Launcher context

export NCCL_DEBUG=INFO
deepspeed --hostfile hosts --bind_cores_to_rank training_script.py --deepspeed $@

Docker context
no docker

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 by reproducing the provided training script with the shown ZeRO stage-2 configuration and torch profiler. Inspect the overlap_comm communication scheduling path and compare the profile with the expected backward/NCCL overlap. Done means identifying the cause and confirming either the expected overlap or a documented limitation.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.