modal/train: VGROUT_ATTN attn-impl override (NOT a fix for the modal hang)

Adds env override VGROUT_ATTN (default flash_attention_2, so local behavior is
unchanged; app.py sets sdpa on Modal). Tested to isolate the Modal generate()
deadlock: it hangs at the first generate under BOTH flash_attention_2 and sdpa,
so the hang is NOT the attention backend -- it's in the generation loop, suspect
the cache-frozen image's transformers-main commit differing from local's working
5.8.0.dev0. Diagnosis + fix path in task #212. Local n=3 runs proceed meanwhile.

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>

modal/app.py

@@ -134,6 +134,9 @@ def _run_train(argv: list[str]) -> dict:
         "HF_HOME": f"{CACHE}/hf",
         "HF_HUB_DISABLE_PROGRESS_BARS": "1",
         "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
+        # The Dao flash-attn 2.8.3 wheel deadlocks generate() on Modal's H100/A100;
+        # sdpa is exact attention (identical outputs, only speed/mem differ).
+        "VGROUT_ATTN": "sdpa",
     }
     runs_before = set(Path(f"{CACHE}/out/runs").glob("*")) if Path(f"{CACHE}/out/runs").exists() else set()
 

src/vgrout/train.py

@@ -446,12 +446,15 @@ def main(cfg: Config) -> int:
 
     # ── model + tokenizer ──
     # CPU smoke: fp32 + sdpa (flash-attn2 is CUDA-only, CPU bf16 is patchy).
-    # GPU: bf16 + flash_attention_2.
+    # GPU: bf16 + flash_attention_2, override via VGROUT_ATTN (e.g. sdpa on Modal,
+    # whose Dao flash-attn wheel hangs generate() on H100/A100 -- sdpa is exact
+    # attention, identical outputs, only speed/memory differ).
     cpu = device.type == "cpu"
+    attn_impl = "sdpa" if cpu else os.environ.get("VGROUT_ATTN", "flash_attention_2")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         dtype=torch.float32 if cpu else torch.bfloat16,
-        attn_implementation="sdpa" if cpu else "flash_attention_2",
+        attn_implementation=attn_impl,
     ).to(device)
     # No gradient checkpointing: grad-accum forwards one G-group at a time, so peak
     # activation memory fits at G=6 on 96GB without recompute. δS is a leaf inside