Invalid gradient at index 0 with FSDP ( gpt-model)

whatdhack: <p>Getting the following with a HF  “gpt2-medium” model with 2 mpi ranks.  Looks like gradients are not sharded  - 25731584 is half of 51463168 ? How do I go about fixing this ?</p>
<p>Function SplitWithSizesBackward0 returned an invalid gradient at index 0 - got [51463168] but expected shape compatible with [25731584]</p>
<p>PyTorch used is 2.0.1, cuda is 11.8.</p>

whatdhack: <p>Appears to be coming from .</p>
<p>here</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/237773132d7f9ee474aff444906c133c49cb253b/torch/csrc/autograd/input_metadata.cpp#L145">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/237773132d7f9ee474aff444906c133c49cb253b/torch/csrc/autograd/input_metadata.cpp#L145" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/237773132d7f9ee474aff444906c133c49cb253b/torch/csrc/autograd/input_metadata.cpp#L145" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/237773132d7f9ee474aff444906c133c49cb253b/torch/csrc/autograd/input_metadata.cpp#L145</a></h4>



    <pre class="onebox"><code class="lang-cpp">
      <ol class="start lines" start="135" style="counter-reset: li-counter 134 ;">
          <li>at::Tensor InputMetadata::reduce_grad(at::Tensor&amp; grad) const {</li>
          <li>  // reduce_grad should only be called if is_expandable_to_shape returns true.</li>
          <li>  TORCH_INTERNAL_ASSERT(maybe_expandable_to(grad));</li>
          <li>  return at::sum_to(std::move(grad), shape_as_dim_vector());</li>
          <li>}</li>
          <li></li>
          <li>std::stringstream InputMetadata::incompatible_shape_error_message(</li>
          <li>    const size_t index,</li>
          <li>    const at::Tensor&amp; grad) const {</li>
          <li>  std::stringstream ss{};</li>
          <li class="selected">  ss &lt;&lt; "invalid gradient at index " &lt;&lt; index &lt;&lt; " - got ";</li>
          <li>  if (::torch::autograd::is_cpp_nested_tensor(grad)) {</li>
          <li>    ss &lt;&lt; grad._nested_tensor_size();</li>
          <li>  } else {</li>
          <li>    ss &lt;&lt; grad.sym_sizes();</li>
          <li>  }</li>
          <li>  ss &lt;&lt; " but expected shape compatible with ";</li>
          <li>  if (is_cpp_nested_tensor()) {</li>
          <li>    ss &lt;&lt; shape_as_tensor();</li>
          <li>  } else {</li>
          <li>    ss &lt;&lt; shape_as_dim_vector();</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>through here</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/e9ebda29d87ce0916ab08c06ab26fd3766a870e5/torch/autograd/__init__.py#L200">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/e9ebda29d87ce0916ab08c06ab26fd3766a870e5/torch/autograd/__init__.py#L200" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/e9ebda29d87ce0916ab08c06ab26fd3766a870e5/torch/autograd/__init__.py#L200" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/e9ebda29d87ce0916ab08c06ab26fd3766a870e5/torch/autograd/__init__.py#L200</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="190" style="counter-reset: li-counter 189 ;">
          <li>        tuple(inputs) if inputs is not None else tuple()</li>
          <li></li>
          <li>    grad_tensors_ = _tensor_or_tensors_to_tuple(grad_tensors, len(tensors))</li>
          <li>    grad_tensors_ = _make_grads(tensors, grad_tensors_, is_grads_batched=False)</li>
          <li>    if retain_graph is None:</li>
          <li>        retain_graph = create_graph</li>
          <li></li>
          <li>    # The reason we repeat same the comment below is that</li>
          <li>    # some Python versions print out the first line of a multi-line function</li>
          <li>    # calls in the traceback and some print out the last line</li>
          <li class="selected">    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass</li>
          <li>        tensors, grad_tensors_, retain_graph, create_graph, inputs,</li>
          <li>        allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward pass</li>
          <li></li>
          <li>def grad(</li>
          <li>    outputs: _TensorOrTensors,</li>
          <li>    inputs: _TensorOrTensors,</li>
          <li>    grad_outputs: Optional[_TensorOrTensors] = None,</li>
          <li>    retain_graph: Optional[bool] = None,</li>
          <li>    create_graph: bool = False,</li>
          <li>    only_inputs: bool = True,</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


whatdhack: <p>This is caused by the HF Accelerate package.  HF Accelerate requires setting ACCELERATE_USE_FSDP=true  for FSDP .    Without that I think the Pyorch expects something, but the  model offers something else.</p>