deepspeedai / deepspeedai/DeepSpeed
no overlap while performing stage-2 zero dp training
Nobody has claimed this yet.
- 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:
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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