Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
Megatron-LM
Commits
c1e4526b
Commit
c1e4526b
authored
Jan 13, 2022
by
Lawrence McAfee
Browse files
renamed 'free' -> 'deallocate'
parent
d10f81c5
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
7 additions
and
7 deletions
+7
-7
megatron/schedules.py
megatron/schedules.py
+7
-7
No files found.
megatron/schedules.py
View file @
c1e4526b
...
@@ -76,8 +76,8 @@ def get_forward_backward_func():
...
@@ -76,8 +76,8 @@ def get_forward_backward_func():
# )
# )
# # <<<
# # <<<
# <<<
# <<<
def
fre
e_output_tensor
(
out
):
def
deallocat
e_output_tensor
(
out
):
'''Pseudo-
fre
e (i.e., set to scalar) the output tensor's '.data' field.
'''Pseudo-
deallocat
e (i.e., set to scalar) the output tensor's '.data' field.
This method should be called right after the output tensor has been
This method should be called right after the output tensor has been
sent to the next pipeline stage. At this point, the output tensor is
sent to the next pipeline stage. At this point, the output tensor is
...
@@ -96,7 +96,7 @@ def free_output_tensor(out):
...
@@ -96,7 +96,7 @@ def free_output_tensor(out):
def
custom_backward
(
output
,
grad_output
):
def
custom_backward
(
output
,
grad_output
):
'''Directly call C++ autograd engine.
'''Directly call C++ autograd engine.
To make the '
fre
e_output_tensor' (above) optimization work, the C++
To make the '
deallocat
e_output_tensor' (above) optimization work, the C++
autograd engine must be called directly, bypassing Pytorch's
autograd engine must be called directly, bypassing Pytorch's
torch.autograd.backward. Pytorch's 'backward' checks that the output and
torch.autograd.backward. Pytorch's 'backward' checks that the output and
grad have the same shape, while C++'s 'backward' does not.
grad have the same shape, while C++'s 'backward' does not.
...
@@ -428,7 +428,7 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat
...
@@ -428,7 +428,7 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat
# >>>
# >>>
pax
({
"output_tensor"
:
output_tensor
})
pax
({
"output_tensor"
:
output_tensor
})
# <<<
# <<<
fre
e_output_tensor
(
output_tensor
)
deallocat
e_output_tensor
(
output_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
):
...
@@ -492,7 +492,7 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat
...
@@ -492,7 +492,7 @@ def forward_backward_pipelining_with_interleaving(forward_step_func, data_iterat
output_tensor
,
input_tensor_grad
,
output_tensor
,
input_tensor_grad
,
recv_prev
=
recv_prev
,
recv_next
=
recv_next
,
recv_prev
=
recv_prev
,
recv_next
=
recv_next
,
tensor_shape
=
tensor_shape
,
timers
=
timers
)
tensor_shape
=
tensor_shape
,
timers
=
timers
)
fre
e_output_tensor
(
output_tensor
)
deallocat
e_output_tensor
(
output_tensor
)
# Put input_tensor and output_tensor_grad in data structures in the
# Put input_tensor and output_tensor_grad in data structures in the
# right location.
# right location.
...
@@ -668,7 +668,7 @@ def forward_backward_pipelining_without_interleaving(forward_step_func, data_ite
...
@@ -668,7 +668,7 @@ def forward_backward_pipelining_without_interleaving(forward_step_func, data_ite
if
not
forward_only
:
if
not
forward_only
:
input_tensors
.
append
(
input_tensor
)
input_tensors
.
append
(
input_tensor
)
output_tensors
.
append
(
output_tensor
)
output_tensors
.
append
(
output_tensor
)
fre
e_output_tensor
(
output_tensor
[
0
])
deallocat
e_output_tensor
(
output_tensor
[
0
])
# Before running 1F1B, need to receive first forward tensor.
# Before running 1F1B, need to receive first forward tensor.
# If all microbatches are run in warmup / cooldown phase, then no need to
# If all microbatches are run in warmup / cooldown phase, then no need to
...
@@ -697,7 +697,7 @@ def forward_backward_pipelining_without_interleaving(forward_step_func, data_ite
...
@@ -697,7 +697,7 @@ def forward_backward_pipelining_without_interleaving(forward_step_func, data_ite
# Add input_tensor and output_tensor to end of list.
# Add input_tensor and output_tensor to end of list.
input_tensors
.
append
(
input_tensor
)
input_tensors
.
append
(
input_tensor
)
output_tensors
.
append
(
output_tensor
)
output_tensors
.
append
(
output_tensor
)
fre
e_output_tensor
(
output_tensor
[
0
])
deallocat
e_output_tensor
(
output_tensor
[
0
])
# 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
# the backward pass.
# the backward pass.
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment