deepspeedai / deepspeedai/DeepSpeed

Grad parameters are None with bigger autograd graphs

Open
#4,941 5 comments 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

I'm training a model that involves using multiple LoRA adapters (3 different sets of adapters to be precise). For each input (which in my case is an image) I have to pass one version through one set of LoRA parameters, another version of the image through another set of LoRA and the final version of the image through the final set of LoRA adapters. This consumes a lot of memory in the form of the autograd graph, I believe.

What is happening is that when I use the 3 sets of LoRA adapters I get errors in DeepSpeed claiming the the grad parameters are None. Here is the traceback using zero stage 3:

Traceback (most recent call last):
  File "/cfs/home/u021543/pheye_llavar_accelerate.py", line 68, in <module>
  0%|                                                                                       | 0/16 [00:02<?, ?it/s]
    accelerator.backward(loss)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/accelerate/accelerator.py", line 1958, in backward
Traceback (most recent call last):
  File "/cfs/home/u021543/pheye_llavar_accelerate.py", line 68, in <module>
    accelerator.backward(loss)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/accelerate/accelerator.py", line 1958, in backward
    self.deepspeed_engine_wrapped.backward(loss, **kwargs)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/accelerate/utils/deepspeed.py", line 167, in backward
    self.engine.backward(loss, **kwargs)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
        self.deepspeed_engine_wrapped.backward(loss, **kwargs) 
         ^^^^^^^^^^  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/accelerate/utils/deepspeed.py", line 167, in backward
^^^^^^^^^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 1955, in backward
    self.engine.backward(loss, **kwargs)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 1955, in backward
    self.optimizer.backward(loss, retain_graph=retain_graph)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 2135, in backward
    self.optimizer.backward(loss, retain_graph=retain_graph)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 2135, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/torch/_tensor.py", line 492, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
    torch.autograd.backward(
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/torch/autograd/__init__.py", line 251, in backward
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
    scaled_loss.backward(retain_graph=retain_graph)
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/torch/_tensor.py", line 492, in backward
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 1119, in reduce_partition_and_remove_grads
    torch.autograd.backward(
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^    ^self.reduce_ready_partitions_and_remove_grads(param)^
^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 1119, in reduce_partition_and_remove_grads
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 1409, in reduce_ready_partitions_and_remove_grads
    self.reduce_ready_partitions_and_remove_grads(param)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 1409, in reduce_ready_partitions_and_remove_grads
    self.reduce_independent_p_g_buckets_and_remove_grads(param)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 1156, in reduce_independent_p_g_buckets_and_remove_grads
    self.reduce_independent_p_g_buckets_and_remove_grads(param)    
self.__add_grad_to_ipg_bucket(param)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 1156, in reduce_independent_p_g_buckets_and_remove_grads
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
         self.__add_grad_to_ipg_bucket(param) 
     ^^^^^  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
^^^^^^^^^^^^^^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 1164, in __add_grad_to_ipg_bucket
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^    ^if self.contiguous_gradients and self.elements_in_ipg_bucket + param.grad.numel() <= self.reduce_bucket_size:^
^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 1164, in __add_grad_to_ipg_bucket
                                                                   ^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'numel'
    if self.contiguous_gradients and self.elements_in_ipg_bucket + param.grad.numel() <= self.reduce_bucket_size:
                                                                   ^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'numel'

and here is the traceback using zero stage 2:

Traceback (most recent call last):
  File "/cfs/home/u021543/pheye_llavar_accelerate.py", line 68, in <module>
  0%|                                                                                                                                                                                                                | 0/16 [00:08<?, ?it/s]
    accelerator.backward(loss)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/accelerate/accelerator.py", line 1958, in backward
Traceback (most recent call last):
  File "/cfs/home/u021543/pheye_llavar_accelerate.py", line 68, in <module>
    accelerator.backward(loss)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/accelerate/accelerator.py", line 1958, in backward
    self.deepspeed_engine_wrapped.backward(loss, **kwargs)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/accelerate/utils/deepspeed.py", line 167, in backward
    self.engine.backward(loss, **kwargs)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 1955, in backward
    self.deepspeed_engine_wrapped.backward(loss, **kwargs)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/accelerate/utils/deepspeed.py", line 167, in backward
    self.optimizer.backward(loss, retain_graph=retain_graph)
    self.engine.backward(loss, **kwargs)  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 2019, in backward

  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 1955, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/torch/_tensor.py", line 492, in backward
    self.optimizer.backward(loss, retain_graph=retain_graph)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 2019, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/torch/autograd/__init__.py", line 251, in backward
    torch.autograd.backward(
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 865, in reduce_partition_and_remove_grads
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 865, in reduce_partition_and_remove_grads
    self.reduce_ready_partitions_and_remove_grads(param, i)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1377, in reduce_ready_partitions_and_remove_grads
    self.reduce_ready_partitions_and_remove_grads(param, i)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1377, in reduce_ready_partitions_and_remove_grads
    self.reduce_independent_p_g_buckets_and_remove_grads(param, i)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 911, in reduce_independent_p_g_buckets_and_remove_grads
    self.reduce_independent_p_g_buckets_and_remove_grads(param, i)
  File "/cfs/home/u021543/miniconda3/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 911, in reduce_independent_p_g_buckets_and_remove_grads
    new_grad_tensor.copy_(grad_reduc.view(-1))
                          ^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'view'
    new_grad_tensor.copy_(grad_reduc.view(-1))
                          ^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'view'

My problem is very similar to https://github.com/microsoft/DeepSpeed/issues/700#issue-795318541, because when using zero stage 1 I get no issue, but I really need to use stage 2.

System info

  • OS: openSUSE Leap 15.5
  • One machine with 2 NVIDIA Quadro RTX 6000 (25GB)
  • Python 3.11.4
  • deepspeed==0.12.6
  • torch==2.1.1

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 with the reported entry point in deepspeed/runtime/zero/stage3.py, especially __add_grad_to_ipg_bucket, then compare the corresponding flow in deepspeed/runtime/zero/stage_1_and_2.py. Reproduce the failure using the three-LoRA setup and the provided accelerator.backward(loss) call, checking why param.grad is None. Done means the issue is reduced to a reproducible case or the affected ZeRO paths handle it without the reported AttributeError.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.