[BUG] When using LLaVA with freeze-LM, training text only sample occurs error.

[liveseongho]

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

	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])

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?