# How to apply selective activation checkpointing on \_grouped\_mm

**URL:** <https://discuss.pytorch.org/t/how-to-apply-selective-activation-checkpointing-on-grouped-mm/223226>\
**Category:** torchtitan\
**Created:** [September 20, 2025, 11:02pm UTC](https://discuss.pytorch.org/t/how-to-apply-selective-activation-checkpointing-on-grouped-mm/223226 "2025-09-20T23:02:37Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![go\_brrr](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/go_brrr/32/78269_2.png) [@go\_brrr](https://discuss.pytorch.org/u/go_brrr)\
**Post date:** [September 20, 2025, 11:02pm UTC](https://discuss.pytorch.org/t/how-to-apply-selective-activation-checkpointing-on-grouped-mm/223226/1 "2025-09-20T23:02:37Z")

</div>

From the selective activation checkpointing blog, I understood that we can save compute intensive operations while recompute operations that are memory intensive. How do I add `_grouped_mm` to `ops_to_save` below?

Does directly adding `torch.ops.aten.grouped_mm.default` works or can I directly adding `torch._grouped_mm` does the trick ?

> ops\_to\_save = [
> 
> torch.ops.aten.mm.default,
> 
> torch.ops.aten.bmm.default,
> 
> torch.ops.aten.addmm.default,
> 
> torch.ops.aten.grouped\_mm.default,
> 
> _# Save collectives that produce dynamic-shaped outputs_
> 
> torch.ops.\_c10d\_functional.all\_to\_all\_single.default,
> 
> torch.ops.\_c10d\_functional.reduce\_scatter\_tensor.default,
> 
> _# for low precision training, it’s useful to always save_
> 
> _# the result of max, since the absolute maximum is_
> 
> _# used to compute the scaling factor for quantization._
> 
> torch.ops.aten.max.default,
> 
> ]
> 
> def \_apply\_ac(_model_: torch.nn.Module):
> 
> _from_ torch.utils.checkpoint _import_ CheckpointPolicy, create\_selective\_checkpoint\_contexts
> 
> _from_ torch.distributed.algorithms.\_checkpoint.checkpoint\_wrapper _import_ checkpoint\_wrapper, CheckpointImpl
> 
> _# Create a policy function that returns a CheckpointPolicy_
> 
> def policy\_fn(_ctx_, _op_, \*_args_, \*\*_kwargs_):
> 
> _if_ _op_ in ops\_to\_save:
> 
> _return_ CheckpointPolicy.MUST\_SAVE
> 
> _else_:
> 
> _return_ CheckpointPolicy.PREFER\_RECOMPUTE
> 
> “”“Apply activation checkpointing to the model.”“”
> 
> _for_ layer\_id, transformer\_block _in_ _model_.layers.named\_children():
> 
> transformer\_block = checkpoint\_wrapper(
> 
> transformer\_block,
> 
> _checkpoint\_impl_=CheckpointImpl.NO\_REENTRANT,
> 
> _context\_fn_=partial(create\_selective\_checkpoint\_contexts, policy\_fn),
> 
> )
> 
> _model_.layers.register\_module(layer\_id, transformer\_block)
> 
> logger.info(“Applied selective activation checkpointing to the model”)

I am using moe implementation from torch titan from reference, I wan