Is torch_remat intended to avoid per-op SAC dispatch overhead in host-bound workloads?

[DavidFFFan]

Hi, thanks for open-sourcing this project.

We encountered a performance issue when using PyTorch native Selective Activation Checkpointing (SAC) for a Transformer layer.

Our intended policy was:

• recompute the Transformer layer as a whole;
• save the All-to-All operation and exclude it from recomputation;
• recompute the remaining operations.

Conceptually, the SAC policy looked like this:

def policy_fn(ctx, op, *args, **kwargs):
    if is_all_to_all(op):
        return CheckpointPolicy.MUST_SAVE
    return CheckpointPolicy.PREFER_RECOMPUTE

Unexpectedly, this selective configuration was slower than fully recomputing the entire Transformer layer, even though it avoided repeating the All-to-All operation.

Profiler analysis suggested that the workload became Host-bound. In particular, the per-operation TorchDispatchMode / policy callback path appeared to delay accelerator operator launches. The additional Host dispatch overhead outweighed the compute/communication saved by excluding All-to-All from recomputation.

Is this type of workload one of the main motivations for torch_remat?

Our understanding is that the equivalent configuration would be approximately:

def transformer_layer(x):
    # Other operations recompute by default.
    x = remat.region(
        all_to_all_region,
        "all_to_all",
        recompute=False,
    )(x)
    return remaining_layer(x)

output = remat.checkpoint(
    region_name="transformer_layer",
)(transformer_layer)(input)

Unlike native SAC, this makes the save/recompute decision at explicit callable-region boundaries rather than invoking a Python policy for every dispatched ATen operation.

We would like to confirm:

1. Is avoiding per-ATen-op dispatch overhead a design goal of torch_remat?
2. Is wrapping the whole Transformer layer with remat.checkpoint and marking only the All-to-All region with recompute=False the recommended configuration?
3. Are there any benchmarks comparing the Host overhead of region-based rematerialization against native SAC?
4. Are there recommended region granularities for Host-bound workloads?

[ezyang]

torch_remat avoids TorchDispatchMode on the inside entirely, so this may potentially improve the host-side performance. However, in our production configuration we are running CUDA graphs on top of it, so while remat may be a promising way of improving host side perf, I haven't actually benchmarked for this, and there might be some minor problems we have to fix before we can get it. It is a design goal to avoid TorchDispatchMode for torch_remat, but host perf is more of an incidental thing. If you like, you can ask your agent to generate a representative config demonstrating host boundedness and I am happy to have an agent work on micro-optimizations to improve perf in this case.

Re how to use torch_remat for your case, as you have described your desired policy, yes you would wrap the all_to_all with recompute=False. One region for you, should be relatively cheap.