refactor: update training scripts and environment setup for Qwen3 model

Modified training scripts to transition from Qwen2.5 to Qwen3, including updates to model and adapter paths in `train_rlvr.py` and `train_grpo.py`. Removed legacy environment variables from setup scripts and README.md, introducing new command-line arguments for stub workers and LLM judge control. Enhanced documentation to reflect these changes and improve user guidance for the updated training process.

env_artifacts/README.md

@@ -61,8 +61,6 @@ pip install -e .          # picks up project deps without touching the pinned to
 
 # Per-session exports (stick these in ~/.bashrc if you want):
 export HF_HUB_ENABLE_HF_TRANSFER=1
-export CORP_STUB_WORKERS=1
-export CORP_DISABLE_LLM_JUDGE=1
 export TOKENIZERS_PARALLELISM=false
 export TRANSFORMERS_VERBOSITY=warning
 
@@ -81,17 +79,17 @@ You can pull them locally on a new box with:
 huggingface-cli download Navigam/corp-env-sft-qwen2.5-7b --local-dir outputs/sft_adapter
 ```
 
-Then for the RLVR run:
+Then for a direct 14B RLVR run:
 
 ```bash
 python training/train_rlvr.py \
-  --model Qwen/Qwen2.5-7B-Instruct \
-  --adapter outputs/sft_adapter \
+  --model Qwen/Qwen3-14B-Instruct \
+  --adapter outputs/sft_qwen3_14b \
   --examples data/processed/e1_m1_clean.jsonl,data/processed/h1_seed_clean.jsonl \
-  --output outputs/rlvr_adapter \
+  --output outputs/rlvr_qwen3_14b \
   --rounds 3 --n-samples 8 --max-prompts 128 \
   --stats-file results/runs/rlvr_stats.jsonl \
-  --push-to-hub Navigam/corp-env-rlvr-qwen2.5-7b
+  --push-to-hub <your-user-or-org>/corp-env-rlvr-qwen3-14b
 ```
 
 ## Notes / gotchas
@@ -99,12 +97,10 @@ python training/train_rlvr.py \
 - **First time only on any new box**: `unsloth` writes patched-trainer classes
   into `corp_gym/unsloth_compiled_cache/`. Delete that folder if you ever
   change TRL/unsloth versions to avoid stale compiled patches.
-- **Flash-attn fp32 bug**: in `training/train_grpo.py` (legacy) and
-  `training/train_rlvr.py` we monkey-patch `flash_attn_func` in both
-  `flash_attn.flash_attn_interface` and `unsloth.utils.attention_dispatch` to
-  auto-cast Q/K/V to bf16. Without the dispatcher-level patch, Unsloth's copy
-  of the symbol bypasses ours and FA2 rejects fp32 tensors during the no-grad
-  reference logprob forward.
+- **Flash-attn fp32 workaround**: this is now **legacy-only** for
+  `training/train_grpo.py`. It is disabled by default and can be enabled only
+  when needed via `CORP_ENABLE_FA2_BF16_PATCH=1`.
+- `training/train_rlvr.py` now runs without this monkey patch on modern stacks.
 - **max_prompt_length filter**: both legacy GRPO and RLVR scripts tokenise every
   prompt up-front and drop rows whose chat-template-encoded length exceeds
   `0.9 * max_prompt_length` (long H1 trajectories otherwise produce a causal

env_artifacts/setup_lightning_h100.sh

@@ -71,8 +71,6 @@ log "done. Export env vars in your shell (rc file or per-session):"
 cat <<'EOF'
 
   export HF_HUB_ENABLE_HF_TRANSFER=1
-  export CORP_STUB_WORKERS=1
-  export CORP_DISABLE_LLM_JUDGE=1
   export TOKENIZERS_PARALLELISM=false
   export TRANSFORMERS_VERBOSITY=warning
 

training/train_grpo.py

@@ -48,8 +48,17 @@ from typing import Any, Dict, List
 
 import torch
 
-try:
-    import flash_attn.flash_attn_interface as _fa_mod
+def _maybe_patch_flash_attn_bf16() -> None:
+    """Legacy workaround for older Unsloth/FA2 stacks.
+
+    Disabled by default; enable only if your GRPO run hits fp32 FA2 dtype errors.
+    """
+    if os.environ.get("CORP_ENABLE_FA2_BF16_PATCH", "").strip() != "1":
+        return
+    try:
+        import flash_attn.flash_attn_interface as _fa_mod
+    except ImportError:
+        return
 
     _orig_fa_func = _fa_mod.flash_attn_func
     _orig_fa_var = getattr(_fa_mod, "flash_attn_varlen_func", None)
@@ -67,16 +76,14 @@ try:
     _fa_mod.flash_attn_func = _fa_func_bf16
     if _orig_fa_var is not None:
         _fa_mod.flash_attn_varlen_func = _fa_varlen_bf16
-
     try:
         import unsloth.utils.attention_dispatch as _ad_mod
+
         _ad_mod.flash_attn_func = _fa_func_bf16
         if _orig_fa_var is not None and hasattr(_ad_mod, "flash_attn_varlen_func"):
             _ad_mod.flash_attn_varlen_func = _fa_varlen_bf16
     except Exception:
         pass
-except ImportError:
-    pass
 
 ROOT = Path(__file__).resolve().parents[1]
 if str(ROOT) not in sys.path:
@@ -343,10 +350,23 @@ def main() -> None:
         default=1,
         help="Filter training examples to traces with at least this many log_reasoning actions.",
     )
+    parser.add_argument(
+        "--use-stub-workers",
+        action="store_true",
+        help="Use deterministic stub workers (default: real workers).",
+    )
+    parser.add_argument(
+        "--disable-llm-judge",
+        action="store_true",
+        help="Disable LLM judge scoring for deterministic verifier-only runs.",
+    )
     args = parser.parse_args()
 
-    os.environ.setdefault("CORP_STUB_WORKERS", "1")
-    os.environ.setdefault("CORP_DISABLE_LLM_JUDGE", "1")
+    if args.use_stub_workers:
+        os.environ["CORP_STUB_WORKERS"] = "1"
+    if args.disable_llm_judge:
+        os.environ["CORP_DISABLE_LLM_JUDGE"] = "1"
+    _maybe_patch_flash_attn_bf16()
 
     try:
         from unsloth import FastLanguageModel, PatchFastRL

training/train_rlvr.py

@@ -22,12 +22,12 @@ completions, and the advantage-variance collapse in GRPO no longer applies.
 Run (after SFT)::
 
     python training/train_rlvr.py \\
-      --model Qwen/Qwen2.5-7B-Instruct \\
-      --adapter outputs/sft_adapter \\
+      --model Qwen/Qwen3-14B-Instruct \\
+      --adapter outputs/sft_qwen3_14b \\
       --examples data/processed/e1_m1_clean.jsonl,data/processed/h1_seed_clean.jsonl \\
-      --output outputs/rlvr_adapter \\
+      --output outputs/rlvr_qwen3_14b \\
       --rounds 3 --n-samples 8 --max-prompts 64 \\
-      --push-to-hub Navigam/corp-env-rlvr-qwen2.5-7b
+      --push-to-hub your-org/corp-env-rlvr-qwen3-14b
 """
 
 from __future__ import annotations
@@ -45,36 +45,6 @@ from typing import Any, Dict, List, Tuple
 
 import torch
 
-try:
-    import flash_attn.flash_attn_interface as _fa_mod
-
-    _orig_fa_func = _fa_mod.flash_attn_func
-    _orig_fa_var = getattr(_fa_mod, "flash_attn_varlen_func", None)
-
-    def _fa_func_bf16(q, k, v, *args, **kwargs):
-        if q.dtype not in (torch.bfloat16, torch.float16):
-            q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)
-        return _orig_fa_func(q, k, v, *args, **kwargs)
-
-    def _fa_varlen_bf16(q, k, v, *args, **kwargs):
-        if q.dtype not in (torch.bfloat16, torch.float16):
-            q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)
-        return _orig_fa_var(q, k, v, *args, **kwargs)
-
-    _fa_mod.flash_attn_func = _fa_func_bf16
-    if _orig_fa_var is not None:
-        _fa_mod.flash_attn_varlen_func = _fa_varlen_bf16
-
-    try:
-        import unsloth.utils.attention_dispatch as _ad_mod
-        _ad_mod.flash_attn_func = _fa_func_bf16
-        if _orig_fa_var is not None and hasattr(_ad_mod, "flash_attn_varlen_func"):
-            _ad_mod.flash_attn_varlen_func = _fa_varlen_bf16
-    except Exception:
-        pass
-except ImportError:
-    pass
-
 ROOT = Path(__file__).resolve().parents[1]
 if str(ROOT) not in sys.path:
     sys.path.insert(0, str(ROOT))
@@ -345,8 +315,8 @@ def maybe_push_to_hub(output_dir: str, repo_id: str) -> None:
 
 def main() -> None:
     parser = argparse.ArgumentParser(description="RLVR / Rejection-Sampling FT for CORP-ENV.")
-    parser.add_argument("--model", default="Qwen/Qwen2.5-7B-Instruct")
-    parser.add_argument("--adapter", default="outputs/sft_adapter",
+    parser.add_argument("--model", default="Qwen/Qwen3-14B-Instruct")
+    parser.add_argument("--adapter", default="outputs/sft_qwen3_14b",
                         help="Start from this LoRA adapter (usually SFT).")
     parser.add_argument(
         "--examples",
@@ -356,7 +326,7 @@ def main() -> None:
         "--tasks",
         default="e1_launch_readiness,m1_budget_reallocation,h1_acquisition_defence",
     )
-    parser.add_argument("--output", default="outputs/rlvr_adapter")
+    parser.add_argument("--output", default="outputs/rlvr_qwen3_14b")
     parser.add_argument("--rounds", type=int, default=3,
                         help="Outer rollout -> filter -> SFT iterations.")
     parser.add_argument("--n-samples", type=int, default=8,
@@ -396,10 +366,22 @@ def main() -> None:
         default="",
         help="Optional JSONL to append per-round rollout stats to.",
     )
+    parser.add_argument(
+        "--use-stub-workers",
+        action="store_true",
+        help="Use deterministic stub workers (default: real workers).",
+    )
+    parser.add_argument(
+        "--disable-llm-judge",
+        action="store_true",
+        help="Disable LLM judge scoring for deterministic verifier-only runs.",
+    )
     args = parser.parse_args()
 
-    os.environ.setdefault("CORP_STUB_WORKERS", "1")
-    os.environ.setdefault("CORP_DISABLE_LLM_JUDGE", "1")
+    if args.use_stub_workers:
+        os.environ["CORP_STUB_WORKERS"] = "1"
+    if args.disable_llm_judge:
+        os.environ["CORP_DISABLE_LLM_JUDGE"] = "1"
 
     random.seed(args.seed)
     torch.manual_seed(args.seed)