openenv/tbench2: DeepSeek-V4-Flash launcher variant#1662
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new experimental launcher script, run-openenv-tbench2-dsv4.py, designed for training DeepSeek-V4-Flash on Terminal-Bench-2. Feedback on the implementation focuses on enhancing FP8 training stability and hardware compatibility. Specifically, it is recommended to define a Transformer Engine precision configuration to keep the DSA indexer in BF16 (preventing NaNs), dynamically adjust FP8 block scaling based on GPU capability (Blackwell vs. Hopper), and add support for the verified 32-GPU parallel configuration on GB300 clusters.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| SCRIPT_DIR = Path(__file__).resolve().parent | ||
|
|
||
|
|
There was a problem hiding this comment.
Define the _DSV4_TE_PRECISION_CONFIG constant at the module level. This configuration is required to keep the DSA indexer weights_proj in BF16 during FP8 training, preventing numerical instability and NaNs.
SCRIPT_DIR = Path(__file__).resolve().parent
_DSV4_TE_PRECISION_CONFIG = """
configs:
bf16:
transformer_engine_config_type: "TEQuantizationParams"
training_recipe: {}
matchers:
dsa_indexer_weights_proj_bf16:
type: "glob"
enabled: true
pattern: "*.self_attention.indexer.linear_weights_proj"
config: "bf16"
""".strip()| if total_gpus == 64: # 8 nodes x 8 GPUs (H200) | ||
| return ( | ||
| "--tensor-model-parallel-size 8 " | ||
| "--sequence-parallel " | ||
| "--pipeline-model-parallel-size 8 " | ||
| "--decoder-first-pipeline-num-layers 4 " | ||
| "--decoder-last-pipeline-num-layers 3 " | ||
| "--context-parallel-size 1 " | ||
| "--expert-model-parallel-size 8 " | ||
| "--expert-tensor-parallel-size 1 " | ||
| ) |
There was a problem hiding this comment.
The docstring and run_deepseek_v4.py mention GB300 (8 nodes x 4 GPUs) as a verified profile for DeepSeek-V4-Flash. However, _parallel_config currently raises a NotImplementedError for this configuration. Let's add the verified 32-GPU parallel configuration to prevent runtime errors on GB300 clusters.
if total_gpus == 32 and num_gpus_per_node == 4: # 8 nodes x 4 GPUs (GB300)
return (
"--tensor-model-parallel-size 2 "
"--sequence-parallel "
"--pipeline-model-parallel-size 8 "
"--decoder-first-pipeline-num-layers 4 "
"--decoder-last-pipeline-num-layers 3 "
"--context-parallel-size 2 "
"--expert-model-parallel-size 4 "
"--expert-tensor-parallel-size 1 "
)
if total_gpus == 64: # 8 nodes x 8 GPUs (H200)
return (
"--tensor-model-parallel-size 8 "
"--sequence-parallel "
"--pipeline-model-parallel-size 8 "
"--decoder-first-pipeline-num-layers 4 "
"--decoder-last-pipeline-num-layers 3 "
"--context-parallel-size 1 "
"--expert-model-parallel-size 8 "
"--expert-tensor-parallel-size 1 "
)| if args.fp8_training: | ||
| misc_args += "--transformer-impl transformer_engine --bf16 --fp8-format e4m3 --fp8-recipe blockwise " | ||
| misc_args += """--train-env-vars '{"NVTE_FP8_BLOCK_SCALING_FP32_SCALES":"1"}' """ |
There was a problem hiding this comment.
During FP8 training, we need to:
- Dynamically set
NVTE_FP8_BLOCK_SCALING_FP32_SCALESto"0"on Blackwell GPUs (which require pow2 scales) and"1"on Hopper GPUs. - Pass
--te-precision-config-fileto keep the DSA indexerweights_projin BF16, avoiding numerical instability and NaNs (mirroring the verified setup inrun_deepseek_v4.py).
| if args.fp8_training: | |
| misc_args += "--transformer-impl transformer_engine --bf16 --fp8-format e4m3 --fp8-recipe blockwise " | |
| misc_args += """--train-env-vars '{"NVTE_FP8_BLOCK_SCALING_FP32_SCALES":"1"}' """ | |
| if args.fp8_training: | |
| misc_args += "--transformer-impl transformer_engine --bf16 --fp8-format e4m3 --fp8-recipe blockwise " | |
| import torch | |
| fp32_scales = "0" if (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10) else "1" | |
| misc_args += f"--train-env-vars '{{\"NVTE_FP8_BLOCK_SCALING_FP32_SCALES\":\"{fp32_scales}\"}}' " | |
| misc_args += f"--te-precision-config-file {U.save_to_temp_file(_DSV4_TE_PRECISION_CONFIG, 'yaml')} " |
6f576df to
777fe2a
Compare
38401a2 to
92ec9f5
Compare
The per-episode client opened in run() was never closed on any path (success, timeout, or exception), leaking httpx connections/sessions under concurrent rollout. Close it in a finally so every path releases it.
92ec9f5 to
251bcbd
Compare
The single-node TB2 example only ever drives the env server over --openenv-env-url; the optional Daytona sandbox pool was never exercised here. Remove the pool (adapter class + launch plumbing + launcher fields) and its docstrings so the example carries only the code path it uses.
…tling The WebSocket env server holds an FD per live session + Docker connection and leaks sockets on unclean disconnects, so a long run under the default 1024 soft nofile exhausts it and the accept loop fails every connection with EMFILE, silently throttling rollouts. Document `ulimit -n` in the start step and the leakage note.
251bcbd to
0350392
Compare
4889fe3 to
3e97c8b
Compare
0350392 to
cd7ee21
Compare
1 participant
Summary
Adds the DeepSeek-V4-Flash launcher variant
run-openenv-tbench2-dsv4.pyfor the OpenEnv tbench2 agentic RL adapter. Stacked on #1487 (shi/openenv-miles-adapter), which carries the shared adapter, the GLM-4.7-Flash launcher, and the sharedopenenv_launch_common.pyhelpers.Why a separate PR
The DeepSeek-V4 launcher is not exercised by the GLM-4.7-Flash smoke run that validates #1487. Splitting it keeps #1487 scoped to the adapter plus one validated launcher, and isolates the DeepSeek-V4 serving/training profile so it can be reviewed on its own. The launcher reuses the shared GRPO/optimizer/rollout/agent/W&B flags from
openenv_launch_common.py; it only overrides the DeepSeek-V4 perf/sglang/misc profile and its ownScriptArgsdefaults.Test plan
rollout/raw_rewardand GRPO steps advance.