File tree Expand file tree Collapse file tree
research/flow_grpo/tests/test_suites Expand file tree Collapse file tree Original file line number Diff line number Diff line change @@ -45,13 +45,39 @@ run_flow_grpo_recipe() {
4545 cd " $PROJECT_ROOT "
4646 NEMO_RL_ROOT=$( realpath " $PROJECT_ROOT /../.." )
4747
48- # paddleocr pulls the full opencv wheels, whose cv2 needs these system libs.
49- if ! ldconfig -p | grep -q libxcb.so.1; then
50- apt-get update && apt-get install -y --no-install-recommends libxcb1 libgl1 libglib2.0-0
51- fi
52-
5348 uv sync --locked --inexact --group test
5449
50+ # paddleocr pulls the full opencv wheels, whose cv2 needs these system libs on every node.
51+ uv run --no-sync python - << 'EOF '
52+ import subprocess
53+
54+ import ray
55+ from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
56+
57+ ray.init(address="auto")
58+
59+
60+ @ray.remote(num_cpus=0)
61+ def install_cv2_libs():
62+ subprocess.run(
63+ "ldconfig -p | grep -q libxcb.so.1 || "
64+ "(apt-get update && apt-get install -y --no-install-recommends libxcb1 libgl1 libglib2.0-0)",
65+ shell=True,
66+ check=True,
67+ )
68+
69+
70+ ray.get(
71+ [
72+ install_cv2_libs.options(
73+ scheduling_strategy=NodeAffinitySchedulingStrategy(n["NodeID"], soft=False)
74+ ).remote()
75+ for n in ray.nodes()
76+ if n["Alive"]
77+ ]
78+ )
79+ EOF
80+
5581 # The recipe inherits data.{train,val}.prompt_file from the exemplar config;
5682 # generate the OCR prompt files on first use.
5783 if [[ ! -f data/ocr/train_prompts.jsonl ]]; then
Original file line number Diff line number Diff line change @@ -35,9 +35,9 @@ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
3535if [[ $( jq ' to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS ) -ge $MAX_STEPS ]]; then
3636 uv run tests/check_metrics.py $JSON_METRICS \
3737 ' data["train/loss"]["1"] < 7.0' \
38- ' data["train/loss"]["50"] < 0.4 ' \
39- ' data["train/grad_norm"]["50"] < 17.5 ' \
40- ' data["train/grad_norm"]["50"] > 10 .0'
38+ ' data["train/loss"]["50"] < 0.35 ' \
39+ ' data["train/grad_norm"]["50"] < 14.0 ' \
40+ ' data["train/grad_norm"]["50"] > 8 .0'
4141
4242 # Clean up checkpoint directory after successful run to save space.
4343 rm -rf " $CKPT_DIR "
You can’t perform that action at this time.
0 commit comments