Fix activation checkpointing crash by using use_reentrant=False

Switches gradient_checkpointing_enable() to use non-reentrant
checkpointing, which properly preserves dropout RNG state during
recomputation and resolves the SystemError during loss.backward().

Issue: #3774

Author: sekyondaMeta

beginner_source/mosaic_memory_profiling_tutorial.py

@@ -309,7 +309,9 @@ def run_training_ac(
     model = GPT2LMHeadModel.from_pretrained("gpt2")
 
     if activation_checkpointing:
-        model.gradient_checkpointing_enable()
+        model.gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs={"use_reentrant": False}
+        )
         print("Activation checkpointing is ENABLED")
     else:
         print("Activation checkpointing is DISABLED")