Commit 0c8e8cce authored by Lawrence McAfee's avatar Lawrence McAfee
Browse files

removed several 'assert_viewless_tensor' calls.

parent a1fe4805
...@@ -307,7 +307,6 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat ...@@ -307,7 +307,6 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat
model[model_chunk_id], model[model_chunk_id],
input_tensor, losses_reduced) input_tensor, losses_reduced)
output_tensors[model_chunk_id].append(output_tensor) output_tensors[model_chunk_id].append(output_tensor)
assert_viewless_tensor(output_tensor)
# if forward-only, no need to save tensors for a backward pass # if forward-only, no need to save tensors for a backward pass
if forward_only: if forward_only:
...@@ -341,7 +340,6 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat ...@@ -341,7 +340,6 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat
mpu.set_virtual_pipeline_model_parallel_rank(0) mpu.set_virtual_pipeline_model_parallel_rank(0)
input_tensors[0].append( input_tensors[0].append(
p2p_communication.recv_forward(tensor_shape, timers=timers)) p2p_communication.recv_forward(tensor_shape, timers=timers))
assert_viewless_tensor(input_tensors[0][-1])
for k in range(num_warmup_microbatches): for k in range(num_warmup_microbatches):
output_tensor = forward_step_helper(k) output_tensor = forward_step_helper(k)
...@@ -373,7 +371,6 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat ...@@ -373,7 +371,6 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat
tensor_shape=tensor_shape, tensor_shape=tensor_shape,
timers=timers) timers=timers)
output_tensor_grads[num_model_chunks-1].append(output_tensor_grad) output_tensor_grads[num_model_chunks-1].append(output_tensor_grad)
assert_viewless_tensor(output_tensor_grad)
else: else:
input_tensor = \ input_tensor = \
p2p_communication.send_forward_recv_forward( p2p_communication.send_forward_recv_forward(
...@@ -382,7 +379,6 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat ...@@ -382,7 +379,6 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat
timers=timers) timers=timers)
free_output_tensor(output_tensor, args.deallocate_pipeline_outputs) free_output_tensor(output_tensor, args.deallocate_pipeline_outputs)
input_tensors[next_forward_model_chunk_id].append(input_tensor) input_tensors[next_forward_model_chunk_id].append(input_tensor)
assert_viewless_tensor(input_tensor)
# Run 1F1B in steady state. # Run 1F1B in steady state.
for k in range(num_microbatches_remaining): for k in range(num_microbatches_remaining):
...@@ -452,18 +448,15 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat ...@@ -452,18 +448,15 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat
# right location. # right location.
if recv_prev: if recv_prev:
input_tensors[next_forward_model_chunk_id].append(input_tensor) input_tensors[next_forward_model_chunk_id].append(input_tensor)
assert_viewless_tensor(input_tensor)
if recv_next: if recv_next:
output_tensor_grads[next_backward_model_chunk_id].append( output_tensor_grads[next_backward_model_chunk_id].append(
output_tensor_grad) output_tensor_grad)
assert_viewless_tensor(output_tensor_grad)
# Run cooldown backward passes (flush out pipeline). # Run cooldown backward passes (flush out pipeline).
if not forward_only: if not forward_only:
if all_warmup_microbatches: if all_warmup_microbatches:
output_tensor_grads[num_model_chunks-1].append( output_tensor_grads[num_model_chunks-1].append(
p2p_communication.recv_backward(tensor_shape, timers=timers)) p2p_communication.recv_backward(tensor_shape, timers=timers))
assert_viewless_tensor(output_tensor_grads[num_model_chunks-1][-1])
for k in range(num_microbatches_remaining, num_microbatches): for k in range(num_microbatches_remaining, num_microbatches):
input_tensor_grad = backward_step_helper(k) input_tensor_grad = backward_step_helper(k)
next_backward_model_chunk_id = get_model_chunk_id(k+1, forward=False) next_backward_model_chunk_id = get_model_chunk_id(k+1, forward=False)
...@@ -478,7 +471,6 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat ...@@ -478,7 +471,6 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat
input_tensor_grad, recv_next=recv_next, input_tensor_grad, recv_next=recv_next,
tensor_shape=tensor_shape, tensor_shape=tensor_shape,
timers=timers)) timers=timers))
assert_viewless_tensor(output_tensor_grads[next_backward_model_chunk_id][-1])
return losses_reduced return losses_reduced
...@@ -624,8 +616,8 @@ def forward_backward_pipelining_without_interleaving(forward_step_func, data_ite ...@@ -624,8 +616,8 @@ def forward_backward_pipelining_without_interleaving(forward_step_func, data_ite
send_forward(output_tensor, send_tensor_shapes, timers=timers) send_forward(output_tensor, send_tensor_shapes, timers=timers)
if not forward_only: if not forward_only:
input_tensors.append(mpu.assert_viewless_tensor(input_tensor)) input_tensors.append(input_tensor)
output_tensors.append(mpu.assert_viewless_tensor(output_tensor)) output_tensors.append(output_tensor)
free_output_tensor(output_tensor, args.deallocate_pipeline_outputs) free_output_tensor(output_tensor, args.deallocate_pipeline_outputs)
# Before running 1F1B, need to receive first forward tensor. # Before running 1F1B, need to receive first forward tensor.
...@@ -653,8 +645,8 @@ def forward_backward_pipelining_without_interleaving(forward_step_func, data_ite ...@@ -653,8 +645,8 @@ def forward_backward_pipelining_without_interleaving(forward_step_func, data_ite
timers=timers) timers=timers)
# Add input_tensor and output_tensor to end of list. # Add input_tensor and output_tensor to end of list.
input_tensors.append(mpu.assert_viewless_tensor(input_tensor)) input_tensors.append(input_tensor)
output_tensors.append(mpu.assert_viewless_tensor(output_tensor)) output_tensors.append(output_tensor)
free_output_tensor(output_tensor, args.deallocate_pipeline_outputs) free_output_tensor(output_tensor, args.deallocate_pipeline_outputs)
# Pop input_tensor and output_tensor from the start of the list for # Pop input_tensor and output_tensor from the start of the list for
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment