[BUG] When using LLaVA with freeze-LM, training text only sample occurs error.
Opened this issue · 0 comments
liveseongho commented
Describe the bug
When using LLaVA with freeze-LM, training text only sample occurs error.
To Reproduce
Base code use VQASample for training LLaVA. But, if you use TextSample or samples with text only, it will occur following error.
Stack trace/logs
File "/workspace/megatron/core/pipeline_parallel/schedules.py", line 365, in backward_step
custom_backward(output_tensor[0], output_tensor_grad[0])
File "/workspace/megatron/core/pipeline_parallel/schedules.py", line 149, in custom_backward
Variable._execution_engine.run_backward(
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
Expected behavior
I think output_tensor[0]
always has grad_fn, but it doesn't have grad_fn if any of sample uses parameters with grad.
Environment (please complete the following information):
- Megatron-LM commit ID : 1115e06
Proposed fix
Fix following code
Megatron-LM/megatron/core/pipeline_parallel/schedules.py
Lines 364 to 367 in bd677bf
as follows
if output_tensor[0].requires_grad:
if config.deallocate_pipeline_outputs:
custom_backward(output_tensor[0], output_tensor_grad[0])
else:
torch.autograd.backward(output_tensor[0], grad_tensors=output_tensor_grad[0])
Additional context
@trintamaki could you help to fix this issue?