Skip to content

Commit b715555

Browse files
authored
Remove outdated comments in autodiff (#2442)
1 parent 6ad54ba commit b715555

1 file changed

Lines changed: 1 addition & 6 deletions

File tree

thunder/transforms/autodiff.py

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -461,18 +461,13 @@ def _group_get_grad_bsyms(trace):
461461

462462

463463
def split_into_forward_and_backward(joint_trace: TraceCtx):
464-
"""split a joint trace for forward and backward into separate ones, including recomputation (aka activation checkpointing)"""
464+
"""split a joint trace for forward and backward into separate ones"""
465465

466466
# the joint trace will have the forward computation at the beginning and then the backward computation
467467
# from how it is constructed.
468468
# we split the trace:
469469
# - forward symbols go into forward_part_bsyms
470470
# - all symbols not in the forward go into backward_part_bsyms
471-
# - for recomputation (aka activation checkpointing), we want to insert symbols going into the forward also into the
472-
# backward, but we want to do so "just in time". To this end, we gather the symbols in a dict and later
473-
# insert it when their respective outputs are needed. This is in backward_part_bsyms_recomputed
474-
# The just in time recomputation is a heuristic to save memory mimicking checkpointing: e.g. for a checkpointed
475-
# block, the forward would be recomputed just before computing the gradient.
476471
# the splitting is done in the reverse order of the bound symbols, and works out which bits are needed for the forward
477472
# from there.
478473
forward_part_bsyms = []

0 commit comments

Comments
 (0)