test_activation_checkpoint.py

[githubsgi]

Looks like this code is intended for sanity checking AC . I think it would be good to incorporate the Llama4 like MoE into the code. Would like to hear thoughts on this.

[tianyu-l]

Including mini / fake MoE in AC test sounds reasonable. After all that's why we have the feature / testing in the first place.
cc @soulitzer