用反向重计算减少训练激活占用

1117            if not frame.is_recomputed[gid]:1118                ctx = frame.input_saver.grad_fn1119                args = ctx.get_args(ctx.saved_tensors)11201121                try:1122                    with _recomputation_hook(1123                        weakref.ref(frame), gid1124                    ), torch.autograd.enable_grad():1125                        frame.recompute_fn(*args)1126                except _StopRecomputationError:1127                    pass1128                frame.is_recomputed[gid] = True1129                frame.check_recomputed_tensors_match(gid)

1500    def recompute_fn(*inputs):1501        kwargs, *args = inputs1502        # This will be called later during recomputation. This wrapping enables1503        # the necessary global state to be captured.1504        rng_devices = []1505        if preserve_rng_state and had_device_in_fwd:1506            rng_devices = fwd_devices1507        with torch.random.fork_rng(1508            devices=rng_devices, enabled=preserve_rng_state, device_type=device_type1509        ):1510            if preserve_rng_state:1511                torch.set_rng_state(fwd_cpu_state)1512                if had_device_in_fwd:1513                    set_device_states(fwd_devices, fwd_device_states, device_type=device_type)15141515            device_autocast_ctx = torch.amp.autocast(