Skip to content

feat: support single-GPU GSM8K smoke runs - #1678

Open
CharlesXu-HQ wants to merge 1 commit into
areal-project:mainfrom
CharlesXu-HQ:feature/single-gpu-quickstart
Open

CharlesXu-HQ wants to merge 1 commit into
areal-project:mainfrom
CharlesXu-HQ:feature/single-gpu-quickstart

Conversation

@CharlesXu-HQ

Copy link
Copy Markdown
Contributor

Description

This PR makes the GSM8K GRPO quickstart runnable as a two-step smoke test on a
single 32 GB GPU with one FSDP rank and one SGLang rank.

The single-GPU path uses disk-based weight synchronization so the separated
training and rollout processes do not create a same-GPU NCCL group. It also:

  • avoids starting an unused evaluation rollout when valid_dataset=null;
  • adds reward_max_workers so workstation runs can bound the math reward
    process pool;
  • launches local Python modules with the active interpreter, preserving the
    selected virtual environment; and
  • documents the validated command in both the English and Chinese quickstarts.

This is a compatibility-oriented smoke configuration for validating an
environment before longer experiments. It is not intended as a
throughput-optimized training recipe.

Related Issue

Fixes #310

Type of Change

  • 🐛 Bug fix
  • ✨ New feature
  • 💥 Breaking change
  • 📝 Documentation update
  • ♻️ Refactoring
  • ⚡ Performance improvement
  • ✅ Test coverage improvement

Validation

The documented command was run from main@6feff6df plus this commit on an
NVIDIA GeForce RTX 5090 (32 GB) with Qwen3-0.6B. Both training steps completed,
the process exited with status 0, and no training or GPU process remained after
teardown. Peak observed device usage was 21.83 / 31.36 GiB.

python3 examples/math/gsm8k_rl.py \
    --config examples/math/gsm8k_grpo.yaml \
    scheduler.type=local \
    experiment_name=gsm8k-grpo-single-gpu \
    trial_name=smoke \
    actor.path=Qwen/Qwen3-0.6B \
    actor.backend=fsdp:d1 \
    +actor.attn_impl=sdpa \
    rollout.backend=sglang:d1 \
    +rollout.setup_timeout=900 \
    actor.weight_update_mode=disk \
    enable_offload=false \
    +total_train_steps=2 \
    rollout.max_concurrent_rollouts=2 \
    +reward_max_workers=1 \
    gconfig.n_samples=2 \
    gconfig.max_new_tokens=64 \
    gconfig.max_tokens=512 \
    actor.mb_spec.max_tokens_per_mb=512 \
    actor.optimizer.type=adam_bf16 \
    +actor.optimizer_dtype=bfloat16 \
    train_dataset.batch_size=2 \
    train_dataset.num_workers=0 \
    train_dataset.pin_memory=false \
    +train_dataset.scheduling_spec=null \
    valid_dataset=null \
    sglang.context_length=1024 \
    sglang.mem_fraction_static=0.2 \
    +sglang.max_prefill_tokens=1024 \
    +sglang.attention_backend=torch_native \
    +sglang.sampling_backend=pytorch \
    cluster.n_nodes=1 \
    cluster.n_gpus_per_node=1

Local checks:

  • 32 passed for the focused unit and regression test set;
  • Ruff lint and format checks pass for all changed Python files;
  • mdformat passes for both changed quickstart documents;
  • SPDX license-header checks pass;
  • generated CLI reference files remain unchanged; and
  • git diff --check passes.

pre-commit run --all-files was attempted, but hook initialization could not
finish because the host timed out while fetching the clang-format mirror from
GitHub. The full documentation build was also attempted but could not start
because uv is not installed on the test host. These two checklist items are
left unchecked below.

Checklist

  • I have read the
    Contributing Guide
  • Pre-commit hooks pass (pre-commit run --all-files)
  • Relevant tests pass; new tests added for new functionality
  • Documentation updated (if applicable; built with ./docs/build_all.sh)
  • Branch is up to date with main
  • Self-reviewed via /review-pr command
  • This PR was created by a coding agent via /create-pr
  • This PR is a breaking change

Breaking Change Details (if applicable):

None.

Additional Context

The PyTorch attention and sampling backends favor compatibility on workstation
GPUs. Users can tune the SGLang backends and increase batch size, sequence
length, model size, or training steps after the smoke run succeeds.


Need help? Check the
Contributing Guide
or ask in GitHub Discussions!

Make the GSM8K example usable on a single-GPU workstation without
starting an unused eval engine or an oversized reward process pool.
Use the active Python interpreter for local subprocesses and document
the validated FSDP/SGLang disk-sync recipe.

Key changes:
- Add a configurable reward worker limit to the GSM8K workflow
- Skip eval rollout initialization without a validation dataset
- Document a 32 GB single-GPU smoke configuration

Refs: areal-project#310
Signed-off-by: charlesxu91 <charlesxu.mi@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Does it support reinforcement learning training with a single GPU and sglang?

1 participant