Skip to content

[Pytorch] Fix backward_dw cuda graph order #2376

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ksivaman merged 7 commits into from
Nov 25, 2025
Merged

[Pytorch] Fix backward_dw cuda graph order #2376

Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
bbbb173
fix backward_dw cuda graph order
Wohox Nov 14, 2025
f5d5cc9
Merge branch 'main' into pingtian/fix_wgrad_cuda_graph_order
Wohox Nov 14, 2025
b935500
Merge branch 'main' into pingtian/fix_wgrad_cuda_graph_order
Wohox Nov 17, 2025
79460e1
add validation for num_layers_per_chunk
Wohox Nov 17, 2025
2bddf06
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Nov 17, 2025
a2b4bfc
Merge branch 'main' into pingtian/fix_wgrad_cuda_graph_order
Wohox Nov 19, 2025
0a29da8
Merge branch 'main' into pingtian/fix_wgrad_cuda_graph_order
ksivaman Nov 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 76 additions & 24 deletions transformer_engine/pytorch/graph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import contextlib
import gc
import warnings
from math import ceil
from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union

import torch
Expand Down Expand Up @@ -145,17 +146,22 @@ def _make_graphed_callables(
# values indicate backward passes. Each
# entry in sample_args corresponds to one of the forward
# passes.
num_model_chunks = max(_order)
num_microbatches = len(_order) // num_model_chunks // 2
assert num_model_chunks * num_microbatches * 2 == len(_order)
_order_without_wgrad = []
for c_id in _order:
if ceil(c_id) != c_id:
continue
_order_without_wgrad.append(c_id)
num_model_chunks = max(_order_without_wgrad)
num_microbatches = len(_order_without_wgrad) // num_model_chunks // 2
assert num_model_chunks * num_microbatches * 2 == len(_order_without_wgrad)

# Determine number of layers in each model chunk.
if _num_layers_per_chunk is None:
assert len(sample_args) * 2 >= len(_order) and (
len(sample_args) * 2 % len(_order) == 0
assert len(sample_args) * 2 >= len(_order_without_wgrad) and (
len(sample_args) * 2 % len(_order_without_wgrad) == 0
), (
f"{len(sample_args)} * 2 >= {len(_order)} and {len(sample_args)} * 2 %"
f" {len(_order)} == 0"
f"{len(sample_args)} * 2 >= {len(_order_without_wgrad)} and {len(sample_args)} * 2"
f" % {len(_order_without_wgrad)} == 0"
)
num_layers = len(sample_args) // num_model_chunks // num_microbatches
_num_layers_per_chunk = [num_layers] * num_model_chunks
Expand All @@ -175,7 +181,7 @@ def _make_graphed_callables(
+ f"entries when order input is provided but got {len(callables)}."
)
assert len(sample_args) == total_num_layers * num_microbatches, (
f"Expected {total_num_layers * num_microbatches}"
f"Expected {total_num_layers * num_microbatches} "
+ f"args tuple, but got {len(sample_args)}."
)

Expand Down Expand Up @@ -214,7 +220,7 @@ def _make_graphed_callables(
consumed_sample_q = {}
fwd_idx = [0] * num_model_chunks
for c_id in _order:
m_chunk = abs(c_id) - 1
m_chunk = abs(ceil(c_id)) - 1

if c_id > 0:
sample_start_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + (
Expand All @@ -241,6 +247,8 @@ def _make_graphed_callables(
sample_args[per_callable_fwd_idx] = sample_args[reuse_fwd_idx]
sample_kwargs[per_callable_fwd_idx] = sample_kwargs[reuse_fwd_idx]
fwd_idx[m_chunk] += 1
elif ceil(c_id) != c_id:
continue
else:
num_consumed_samples = min(
len(fwd_sample_qs[m_chunk]), _num_layers_per_chunk[m_chunk]
Expand Down Expand Up @@ -477,9 +485,11 @@ def hook_fn(
fwd_idx = [0] * num_model_chunks
bwd_idx = [0] * num_model_chunks
static_grad_outputs_dict = {}
wgrad_validation_list = [None] * len(_order)
previous_chunk_last_callable_bwd_idx = None
for c_id in _order:
for i, c_id in enumerate(_order):
if c_id > 0:
assert isinstance(c_id, int), "Forward order value must be an integer."
# Capture forward graph for model chunk c_id, microbatch fwd_idx[c_id-1]
m_chunk = c_id - 1
for l_no in range(_num_layers_per_chunk[m_chunk]):
Expand All @@ -499,12 +509,65 @@ def hook_fn(
fwd_idx[m_chunk] += 1
else:
# Capture backward graph for model chunk c_id, microbatch bwd_idx[-c_id-1]
m_chunk = -c_id - 1
m_chunk = -ceil(c_id) - 1
previous_per_callable_bwd_idx = None
for l_no in list(reversed(range(_num_layers_per_chunk[m_chunk]))):
per_callable_bwd_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + (
bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no
)
if ceil(c_id) == c_id and need_bwd_dw_graph[per_callable_bwd_idx]:

@greptile-apps greptile-apps Bot Nov 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: Validation only runs if need_bwd_dw_graph[per_callable_bwd_idx] is true, which depends on the current layer (l_no). If the first layer in reverse order (last layer of the chunk) doesn't need wgrad but other layers do, validation is skipped entirely.

Check if ANY layer in the chunk needs wgrad before running validation, not just the current layer.

@Wohox Wohox Nov 18, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. For delay wgrad case, the layer number is guaranteed to be all 1s, therefore only need to check the first layer
  2. If the current layer doesn't have wgrad, then the check should be skipped. It checks for c_id-0.5 only when there are delay wgrad TE modules.

# Check if bwd graph has corresponding wgrad graph:
# Number of dgrad backward graphs should be equal to number of
# wgrad backward graphs.
# Note: For MCore, the validation rule is more strict (the next backward
# of dgrad graph must be corresponding wgrad graph).
if wgrad_validation_list[i] is None:
same_bwd_c_id_list = [i]
num_wgrad_c_id = 0
for idx in range(i + 1, len(_order)):
if _order[idx] > 0:
continue
if _order[idx] == c_id:
same_bwd_c_id_list.append(idx)
if _order[idx] + 0.5 == c_id:
num_wgrad_c_id += 1
if len(same_bwd_c_id_list) == num_wgrad_c_id:
for same_c_id_idx in same_bwd_c_id_list:
wgrad_validation_list[same_c_id_idx] = True
break
if len(same_bwd_c_id_list) < num_wgrad_c_id:
# It's impossible to have more wgrad than dgrad.
wgrad_validation_list[i] = False
break
if wgrad_validation_list[i] is None:
wgrad_validation_list[i] = False
assert wgrad_validation_list[i], (
f"Number of wgrad graph({num_wgrad_c_id}) doesn't match number "
f"of dgrad graphs ({len(same_bwd_c_id_list)}) for chunk {c_id}."
)
Comment on lines +538 to +561

@greptile-apps greptile-apps Bot Nov 13, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: The validation loop can exit without setting wgrad_validation_list[i], leaving it as None. This happens when the loop completes without finding a match (neither len(same_bwd_c_id_list) == num_wgrad_c_id nor len(same_bwd_c_id_list) < num_wgrad_c_id is true). The assertion on line 542 will then fail with a confusing error since it checks truthiness of None.

Add an else clause after the loop to handle this case:

Suggested change
if wgrad_validation_list[i] is None:
same_bwd_c_id_list = [i]
num_wgrad_c_id = 0
for idx in range(i + 1, len(_order)):
if _order[idx] > 0:
continue
if _order[idx] == c_id:
same_bwd_c_id_list.append(idx)
if _order[idx] + 0.5 == c_id:
num_wgrad_c_id += 1
if len(same_bwd_c_id_list) == num_wgrad_c_id:
for same_c_id_idx in same_bwd_c_id_list:
wgrad_validation_list[same_c_id_idx] = True
break
elif len(same_bwd_c_id_list) < num_wgrad_c_id:
# It's impossible to have more wgrad than dgrad.
wgrad_validation_list[i] = False
break
assert wgrad_validation_list[i], (
f"Number of wgrad graph({num_wgrad_c_id}) doesn't match number "
f"of dgrad graphs ({len(same_bwd_c_id_list)}) for chunk {c_id}."
)
if wgrad_validation_list[i] is None:
same_bwd_c_id_list = [i]
num_wgrad_c_id = 0
for idx in range(i + 1, len(_order)):
if _order[idx] > 0:
continue
if _order[idx] == c_id:
same_bwd_c_id_list.append(idx)
if _order[idx] + 0.5 == c_id:
num_wgrad_c_id += 1
if len(same_bwd_c_id_list) == num_wgrad_c_id:
for same_c_id_idx in same_bwd_c_id_list:
wgrad_validation_list[same_c_id_idx] = True
break
elif len(same_bwd_c_id_list) < num_wgrad_c_id:
# It's impossible to have more wgrad than dgrad.
wgrad_validation_list[i] = False
break
else:
# Loop completed without break - mismatch found
wgrad_validation_list[i] = False
assert wgrad_validation_list[i], (
f"Number of wgrad graph({num_wgrad_c_id}) doesn't match number "
f"of dgrad graphs ({len(same_bwd_c_id_list)}) for chunk {c_id}."
)

@Wohox Wohox Nov 18, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. None means the current check doesn't find the corresponding wgrad chunk_id, and hence should raise assertion.
  2. Assert None and Assert False will both trigger the assertion.

Comment on lines +532 to +561

@greptile-apps greptile-apps Bot Nov 14, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: Validation only runs if the first layer (last in reverse iteration) needs wgrad due to need_bwd_dw_graph[per_callable_bwd_idx] check. If only later layers need wgrad, validation is skipped entirely. The validation logic should check if ANY layer in the chunk needs wgrad, not just per_callable_bwd_idx.

@Wohox Wohox Nov 17, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@buptzyb Do you think we should validate num_layers_per_chunk is a all-1 array when any of the model chunk has need_bwd_dw_graph?

@buptzyb buptzyb Nov 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Make sense. I think this is necessary.

@Wohox Wohox Nov 18, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For delay wgrad case, the layer number is guaranteed to be all 1s, therefore only need to check the first layer

Comment on lines +532 to +561

@greptile-apps greptile-apps Bot Nov 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: validation only runs if need_bwd_dw_graph[per_callable_bwd_idx] is true for the current l_no layer. Since layers are processed in reverse order, if the last layer of a chunk doesn't need wgrad but earlier layers do, validation is skipped entirely. Check if ANY layer in the chunk needs wgrad before running validation

@Wohox Wohox Nov 18, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For delay wgrad case, the layer number is guaranteed to be all 1s, therefore only need to check the first layer

elif ceil(c_id) != c_id:
per_callable_bwd_idx -= _num_layers_per_chunk[m_chunk]

@greptile-apps greptile-apps Bot Nov 14, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: The index adjustment per_callable_bwd_idx -= _num_layers_per_chunk[m_chunk] appears incorrect. This will cause wgrad graphs to be captured from the wrong microbatch or with invalid indices.

When processing a wgrad entry (e.g., -1.5), bwd_idx[m_chunk] has already been incremented by the previous dgrad entry. The correction should likely be at the microbatch level in the index formula, not by subtracting the number of layers. For example, use (bwd_idx[m_chunk] - 1) in the index calculation instead of modifying per_callable_bwd_idx after the fact.

@Wohox Wohox Nov 18, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(bwd_idx[m_chunk] - 1) is mathematically the same as subtracting by _num_layers_per_chunk[m_chunk]

@greptile-apps greptile-apps Bot Nov 18, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're absolutely right! Let me trace through the logic to show why they're mathematically equivalent in this context.

Looking at the code structure:

  1. per_callable_bwd_idx is calculated as:

    per_callable_bwd_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + (
    bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no
)

  2. When we have delay_wgrad_compute enabled, each model chunk has exactly one layer (_num_layers_per_chunk[m_chunk] == 1).

  3. The bwd_idx[m_chunk] gets incremented after processing each dgrad entry.

So when we reach a wgrad entry (where ceil(c_id) != c_id), the bwd_idx[m_chunk] has already been incremented by the previous dgrad processing.

Given that _num_layers_per_chunk[m_chunk] == 1 in the delay_wgrad_compute case:

  • (bwd_idx[m_chunk] - 1) * _num_layers_per_chunk[m_chunk] becomes (bwd_idx[m_chunk] - 1) * 1 = bwd_idx[m_chunk] - 1
  • bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] - _num_layers_per_chunk[m_chunk] becomes bwd_idx[m_chunk] * 1 - 1 = bwd_idx[m_chunk] - 1

You're correct that both approaches yield the same result mathematically. The current implementation works correctly for the delay_wgrad_compute case. My original comment incorrectly identified this as a bug when it's actually functioning as intended.

The key insight is that when delay_wgrad_compute is enabled, _num_layers_per_chunk[m_chunk] is always 1, making the subtraction equivalent to decrementing the microbatch index, which is the correct behavior for accessing the corresponding forward pass tensors for the wgrad computation.

@greptile-apps greptile-apps Bot Nov 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: The index calculation uses bwd_idx[m_chunk] which was already incremented by the previous dgrad entry. For order [1, -1, -1.5], when processing -1.5, bwd_idx[0] is already 1 (from -1), so this calculates index for microbatch 1 then subtracts layers, giving wrong index.

Use (bwd_idx[m_chunk] - 1) in the formula on line 515-516 instead:

Suggested change
per_callable_bwd_idx -= _num_layers_per_chunk[m_chunk]
per_callable_bwd_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + (
(bwd_idx[m_chunk] - 1) * _num_layers_per_chunk[m_chunk] + l_no
) - _num_layers_per_chunk[m_chunk]

@Wohox Wohox Nov 18, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(bwd_idx[m_chunk] - 1) is mathematically the same as subtracting by _num_layers_per_chunk[m_chunk]

@greptile-apps greptile-apps Bot Nov 18, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're absolutely right! Looking at the code more carefully, I can see that per_callable_bwd_idx is calculated using bwd_idx[m_chunk] which gets incremented after each backward pass. So when we're processing the wgrad (float value), bwd_idx[m_chunk] has already been incremented from the previous dgrad pass.

The current line:

per_callable_bwd_idx -= _num_layers_per_chunk[m_chunk]

Is indeed mathematically equivalent to using (bwd_idx[m_chunk] - 1) in the original calculation on lines 529-531, since:

  • Original: (_prefix_num_layers[m_chunk] * num_microbatches) + (bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no)
  • With (bwd_idx[m_chunk] - 1): (_prefix_num_layers[m_chunk] * num_microbatches) + ((bwd_idx[m_chunk] - 1) * _num_layers_per_chunk[m_chunk] + l_no)
  • Current approach: Calculate with bwd_idx[m_chunk], then subtract _num_layers_per_chunk[m_chunk]

Both approaches achieve the same result of using the previous microbatch index. The current implementation is actually more straightforward since it reuses the same calculation pattern and then adjusts the result, rather than modifying the formula itself.

Your suggested fix would work, but the current approach in line 563 is equivalent and arguably cleaner since it maintains consistency with the calculation pattern used elsewhere in the function.

@greptile-apps greptile-apps Bot Nov 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: index calculation is incorrect: bwd_idx[m_chunk] was already incremented by the previous dgrad entry, so this subtracts layers from the wrong microbatch index. When delay_wgrad_compute=True, _num_layers_per_chunk[m_chunk] is always 1, making this per_callable_bwd_idx -= 1, but the real issue is using the wrong base bwd_idx

@Wohox Wohox Nov 18, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

bwd_idx is added by 1, so subtracting by 1 and then multiplies by _num_layers_per_chunk[m_chunk] is mathematically the same as subtracting by _num_layers_per_chunk[m_chunk].

assert is_training, "Only training mode supports backward_dw."
# If no one module needs the backward_dw, the bwd_dw_graph will be empty.
# So skip capturing it. For backward_dw, the order value is c_id - 0.5 to indicate
# the specific order of backward_dw.
assert ceil(c_id) - c_id == 0.5, (
"The order diff of wgrad and dgrad must be 0.5, "
f"get {ceil(c_id) - c_id}."
)
assert need_bwd_dw_graph[
per_callable_bwd_idx
], "No module needs wgrad computation but get float in order"
bwd_dw_graph = bwd_dw_graphs[per_callable_bwd_idx]
with _graph_context_wrapper(bwd_dw_graph, pool=mempool):
for module in visited_te_modules[per_callable_bwd_idx]:
if (
hasattr(module, "need_backward_dw")
and module.need_backward_dw()
):
module.backward_dw()
continue
Comment on lines +562 to +583

@greptile-apps greptile-apps Bot Nov 14, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: The wgrad handling is inside the for l_no loop (line 514), but the continue on line 569 only skips to the next iteration of this inner loop. This means wgrad graphs will be captured multiple times (once for each layer in _num_layers_per_chunk[m_chunk]), which seems incorrect since wgrad should only be captured once per chunk.

Consider using break instead of continue, or restructure the code to handle wgrad entries outside the layer loop.

@Wohox Wohox Nov 18, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only one 1 layer in each model chunk if delay wgrad is enabled, so using continue and break is exactly the same.

Comment on lines +562 to +583

@greptile-apps greptile-apps Bot Nov 14, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: Wgrad graphs will be captured multiple times (once per layer in the chunk) because this code is inside the for l_no loop (line 514). The continue on line 569 only skips to the next layer iteration, not out of the chunk.

For a chunk with 3 layers, the wgrad graph would be captured 3 times with indices calculated as:

  • Layer 0: per_callable_bwd_idx = base + (bwd_idx[m_chunk] * 3 + 0) - 3
  • Layer 1: per_callable_bwd_idx = base + (bwd_idx[m_chunk] * 3 + 1) - 3
  • Layer 2: per_callable_bwd_idx = base + (bwd_idx[m_chunk] * 3 + 2) - 3

The wgrad should only be captured once per chunk. Move this block outside the layer loop or use break instead of continue.

@Wohox Wohox Nov 18, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only one 1 layer in each model chunk if delay wgrad is enabled, so using continue and break is exactly the same.

Comment on lines +562 to +583

@greptile-apps greptile-apps Bot Nov 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: Wgrad capture happens inside the for l_no loop, so it runs once per layer in the chunk. For a 3-layer chunk, wgrad gets captured 3 times with different indices:

  • l_no=2: per_callable_bwd_idx = base + (bwd_idx * 3 + 2) - 3
  • l_no=1: per_callable_bwd_idx = base + (bwd_idx * 3 + 1) - 3
  • l_no=0: per_callable_bwd_idx = base + (bwd_idx * 3 + 0) - 3

Move this entire wgrad block outside the layer loop to capture once per chunk, OR use break instead of continue.

@Wohox Wohox Nov 18, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only one 1 layer in each model chunk if delay wgrad is enabled, so using continue and break is exactly the same.

Comment on lines +562 to +583

@greptile-apps greptile-apps Bot Nov 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: wgrad capture runs inside the for l_no loop (line 528) - once per layer in the chunk. For a chunk with 3 layers, this captures the same wgrad graph 3 times with different indices due to the loop iteration. Move this entire block outside the layer loop

@Wohox Wohox Nov 18, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only one 1 layer in each model chunk if delay wgrad is enabled, so using continue and break is exactly the same.


static_input_surface = per_callable_static_input_surfaces[per_callable_bwd_idx]
static_outputs = per_callable_static_outputs[per_callable_bwd_idx]
bwd_graph = bwd_graphs[per_callable_bwd_idx]
Expand Down Expand Up @@ -537,17 +600,6 @@ def hook_fn(
allow_unused=allow_unused_input,
retain_graph=retain_graph_in_backward,
)
# If no one module needs the backward_dw, the bwd_dw_graph will be empty.
# So skip capturing it.
if need_bwd_dw_graph[per_callable_bwd_idx]:
bwd_dw_graph = bwd_dw_graphs[per_callable_bwd_idx]
with _graph_context_wrapper(bwd_dw_graph, pool=mempool):
for module in visited_te_modules[per_callable_bwd_idx]:
if (
hasattr(module, "need_backward_dw")
and module.need_backward_dw()
):
module.backward_dw()
# Constructs a tuple suitable for returning from Graphed.backward:
# Pads out the actually-needed grads with Nones in gradient slots for inputs
# that don't require grad. I couldn't think of a one-liner for this pattern.
Expand Down Expand Up @@ -596,8 +648,8 @@ def hook_fn(
per_callable_static_grad_inputs[idx]
)
previous_chunk_last_callable_bwd_idx = per_callable_bwd_idx

bwd_idx[m_chunk] += 1
if ceil(c_id) == c_id:
bwd_idx[m_chunk] += 1
else:
# Capture forward graphs
per_callable_static_outputs = []
Expand Down
Loading