File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff line change @@ -461,18 +461,13 @@ def _group_get_grad_bsyms(trace):
461461
462462
463463def 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 = []
You can’t perform that action at this time.
0 commit comments