Cannot execute loss.backward() for training a specific layer

username: Nayoung-Oh
cooked: <p>Hello, my code works when I set requires_grad for other layers, but it suddenly stops working with the specific layer. (Still working when I only enable one GPU)</p>
<p>I am training Transformers using FullyShardedDataParallel (<strong>FSDP)</strong> with <code>ShardingStrategy.HYBRID_SHARD (full-shard), limit_all_gathers=True, use_orig_params=True</code>.</p>
<p>Currently, my loss function is as follows.</p>
<pre><code class="lang-auto">teacher_hidden_states = teacher_output.hidden_states[target_layer_idx+1].detach()
                        student_hidden_states = output.hidden_states[target_layer_idx+1]
loss = F.mse_loss(student_hidden_states, teacher_hidden_states)
</code></pre>
<p>If I add additional CE loss on final prediction, it suddenly works.</p>
<p>The error message I got is as follows.</p>
<pre><code class="lang-auto">[rank1]:     loss.backward()                                                                                                                                                                                       
[rank1]:   File "/root/anaconda3/envs/my_env/lib/python3.10/site-packages/torch/_tensor.py", line 521, in backward                                                                                            
[rank1]:     torch.autograd.backward(                                                                                                                                                                              
[rank1]:   File "/root/anaconda3/envs/my_env/lib/python3.10/site-packages/torch/autograd/__init__.py", line 289, in backward                                                                                  
[rank1]:     _engine_run_backward(                                                                                                                                                                                 
[rank1]:   File "/root/anaconda3/envs/my_env/lib/python3.10/site-packages/torch/autograd/graph.py", line 768, in _engine_run_backward                                                                         
[rank1]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass                                                                                                 
[rank1]:   File "/root/anaconda3/envs/my_env/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 682, in _pre_backward_hook                                                          
[rank1]:     _prefetch_handle(state, handle, _PrefetchMode.BACKWARD)                                                                                                                                               
[rank1]:   File "/root/anaconda3/envs/my_env/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 1222, in _prefetch_handle                                                           
[rank1]:     _unshard(state, handle, state._unshard_stream, state._pre_unshard_stream)                                                                                                                             
[rank1]:   File "/root/anaconda3/envs/my_env/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 300, in _unshard                                                                    
[rank1]:     handle.unshard()                                                                                                                                                                                      
[rank1]:   File "/root/anaconda3/envs/my_env/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 1312, in unshard                                                                       
[rank1]:     self._use_unsharded_flat_param(padded_unsharded_flat_param)                                                                                                                                           
[rank1]:   File "/root/anaconda3/envs/my_env/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 1443, in _use_unsharded_flat_param                                                     
[rank1]:     self._use_unsharded_views(                                                                                                                                                                            
[rank1]:   File "/root/anaconda3/envs/openvla-oft/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context                                                                          
[rank1]:     return func(*args, **kwargs)                                                                                                                                                                          
[rank1]:   File "/root/anaconda3/envs/openvla-oft/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 1925, in _use_unsharded_views
[rank1]:     tensor.data = view                                                                          
[rank1]: AttributeError: 'NoneType' object has no attribute 'data'
</code></pre>
<p>It seems like there is some timing issue….</p>
<p>Has anyone encountered similar issues?</p>
<p>Thanks a lot.</p>

username: Arunprakash-A
cooked: <aside class="quote no-group" data-username="Nayoung-Oh" data-post="1" data-topic="222930">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/nayoung-oh/48/78167_2.png" class="avatar"> Nayoung-Oh:</div>
<blockquote>
<p><code>AttributeError: 'NoneType' object has no attribute 'data'</code></p>
</blockquote>
</aside>
<p>It’s definitely not a timing issue. The error indicates that it tried to access the <code>.data</code> attribute from what was likely a parameter tensor; however, it received a <code>None</code> object instead. Sharing your training loop might help us figure this out</p>