Why `preserve_rng_state=False` in activation checkpointing

[awgu]

Why does torchtitan set preserve_rng_state=False for activation checkpointing? E.g.:

torchtitan/torchtitan/models/llama3/infra/parallelize.py

Line 238 in f4048f8

return ptd_checkpoint_wrapper(module, preserve_rng_state=False)

[tianyu-l]

@soulitzer for any context / suggestions

[soulitzer]

It is surprising to see preserve_rng_state=False. Although you remove the overhead of stashing/restoring RNG state, if any operations use randomness, then you can end up with incorrect gradients.

[tianyu-l]

@soulitzer I see.

I believe Llama is not using randomness -- for any model whose model code doesn't use randomness, it should be fine either way? If a model is using randomness then it's recommended to use preserve_rng_state=True. Is this the right conclusion?

Then I wonder why we have this option at all -- why don't we only support the True case for correctness. Is it because some model doesn't care about incorrect gradients from recomputation, so we can save the overhead?

[soulitzer]

If a model is using randomness then it's recommended to use preserve_rng_state=True. Is this the right conclusion?

Yeah

I believe Llama is not using randomness

Are we using dropout in sdpa btw? (This reminds me that for SAC I might be able to stop treating it as a random op if dropout_p is set to 0)

so we can save the overhead?

We might've checked this at some point to see if the overhead was minor enough to just not care about, I don't think that was the case.

[awgu]

The extra overhead is mainly some extra CPU overhead?

[tianyu-l]

@soulitzer
Llama is not using dropout in sdpa afaik.

We might've checked this at some point to see if the overhead was minor enough to just not care about, I don't think that was the case.

Then why don't we always set it to True?

[soulitzer]

Then why don't we always set it to True?

@tianyu-l Sorry I worded that kind of badly lol, what I mean to say is the overhead didn't seem minor enough to ignore :P

The extra overhead is mainly some extra CPU overhead?

@awgu Not sure, we just looked at the runtime and didn't try to stare too hard at it. Maybe worth checking again 🤔

[sippycoder]

I just experimented with MoE layers in my experiments. Activation checkpointing does not work without preserve_rng_state=True. The router is pretty sensitive to small deviations. There are a number of reasons that can recompute the router differently in backward pass,

• bfloat16/float16 casts
• softmax is extremely sensitive
• CUDA math kernels that are not bit-exact: flash_attn_varlen_func, grouped-GEMM, silu, etc. use CUDA atomics/reductions whose result depends on thread scheduling. They return numerically identical values only up to the last 1-2 umps; two successive launches of the same input generally differ by up to 1 ump or so.

At least for MoE model, preserve_rng_state needs to be True.

torchtitan cleverly skips checkpointing moe.router.gate altogether. So does Megatron. I wouldn't be surprised if preserve_rng_state=True is not enough for checkpointing the router correctly. I will run an experiment and get back.

[tianyu-l]

@sippycoder
Thanks for the note!

torchtitan cleverly skips checkpointing moe.router.gate altogether.

By this you mean torchtitan always recompute moe.router.gate?
For SAC I think it's because of #1380

bfloat16/float16 casts
softmax is extremely sensitive

For this do you think we should consider change moe.router.gate to compute in fp32? Currently I believe it's in bf16, governed by FSDP2 mixed precision behavior.

[sippycoder]

By this you mean torchtitan always recompute moe.router.gate?

Nope. I mean, the router computations are skipped. They try to avoid dealing with the randomness in the re-computations by avoiding the recompute for moe.router.gate and all the fqn's that match same shape as moe.router.gate (a side effect).

For this do you think we should consider change moe.router.gate to compute in fp32? Currently I believe it's in bf16, governed by FSDP2 mixed precision behavior.

It's already done in torch.float32 I believe. Even with that, we can still to avoid the re-computation of routing because it's highly sensitive.

[tianyu-l]

@sippycoder

They try to avoid dealing with the randomness in the re-computations by avoiding the recompute for moe.router.gate and all the fqn's that match same shape as moe.router.gate (a side effect).

What makes you believe that? The code says prefere / force RECOMPUTE moe.router.gate.
https://github.com/pytorch/torchtitan/pull/1380/files#diff-5a556cc02e3e89c77936577320ae0d6c7382ad291ee65d32eea08291dc3618e1R299

It's already done in torch.float32 I believe.

What makes you believe that?

[sippycoder]

I went back and rechecked the code.

What makes you believe that? The code says prefere / force RECOMPUTE moe.router.gate.

Looks like they ARE forcing the recomputing of router gate. I'm surprised, given the sensitivity of it.

moe.router.gate is also not in float32, only the scores coming out of the gate are cast to float32. This is fine as long as recomputations are preserving rng state imo.

[tianyu-l]

@soulitzer

Let's force save of router, instead of force recompute, due to stability concerns.

How do we do this with full AC?

[tianyu-l]

@sippycoder
How is router save / recomputation related to preserve_rng_state=True?
IIUC router.gate is a nn.Linear which doesn't have randomness.