Disable dropout to workaround PyTorch 2.11 checkpoint recomputation bug

Disable dropout (resid_pdrop=0, attn_pdrop=0, embd_pdrop=0) in the
run_training_ac function to avoid SystemError from _VF.dropout returning
NULL during backward recomputation of GPT2Block. Dropout is irrelevant
to the memory profiling purpose of this tutorial.

Issue: #3774

Author: sekyondaMeta

beginner_source/mosaic_memory_profiling_tutorial.py

@@ -306,7 +306,13 @@ def run_training_ac(
 
     # Load model
     print(f"Loading GPT-2 (activation_checkpointing={activation_checkpointing})...")
-    model = GPT2LMHeadModel.from_pretrained("gpt2")
+    # Disable dropout to avoid PyTorch 2.11 checkpoint recomputation bug (#3774).
+    # _VF.dropout returns NULL without setting an exception during backward
+    # recomputation of GPT2Block. Dropout is irrelevant to memory profiling.
+    # Original: model = GPT2LMHeadModel.from_pretrained("gpt2")
+    model = GPT2LMHeadModel.from_pretrained(
+        "gpt2", resid_pdrop=0, attn_pdrop=0, embd_pdrop=0
+    )
 
     if activation_checkpointing:
         model.gradient_checkpointing_enable()