[chenqihong]

Summary

Adds train_dapo.py to the experimental SageMaker backend: the train_grpo.py loop plus two parts of DAPO.

• Dynamic sampling: each step rolls out batches of batch_size prompts and keeps only groups whose rewards differ (an all-equal group has zero advantage), until batch_size groups are kept. Capped by max_num_gen_batches; hitting the cap saves a checkpoint and stops.
• Clip-Higher: clip_low_threshold / clip_high_threshold are sent as loss_fn_config.

Token-level loss and overlong reward shaping are not included: the loss is computed server side, and overlong shaping depends on each agent's reward scale.

Changes

• train_dapo.py: new. train() keeps the structure of train_grpo.py; the differences are the sampling step, loss_fn_config, a per-step sampling: log line (generation batches, filtered fraction, kept groups, empty rollouts), and the final-eval condition. Advantages reuse compute_grpo_advantages.
• config.py: load_config takes an optional config_cls (default Config), so DapoConfig(Config) can add its three fields while train_grpo.py reads the same YAML unchanged.
• config.yaml.example, SETUP.md: document the DAPO fields and how to run train_dapo.py.
• Tests: test_config.py, test_train_dapo.py (sampling logic, and the full train() loop against fake SageMaker/ACR clients).

Two small behavior differences from train_grpo.py:

• Data is loaded before the training session is created, and a dataset smaller than batch_size raises instead of running 0 steps, so bad data never starts a billed session.
• The final eval runs unless an eval already ran at the last step. train_grpo.py uses global_step % eval_every, which skips the final eval when no step ran or the last step was skipped for having no datums. I left train_grpo.py unchanged; happy to fix it in a follow-up.

Testing

• pytest tests/: 655 passed, 23 skipped, 1 failed. The failure is tests/aws_tools/apigateway_tools_test.py::ForwardingTest::test_handles_concurrent_requests, which this PR does not touch; it is flaky locally (passed 1 of 3 reruns). The SageMaker tests (47) all pass; test_train_dapo.py skips when sagemaker.train.training_session is not installed. pre-commit passes on the changed files.
• GSM8K, gpt-oss-20b LoRA, batch 32 x 4, 3,000 prompts, 29 steps. Eval accuracy (200 test questions): 0.69 at step 10, 0.87 at step 20, 0.88 at step 29. This matches an earlier run of the same algorithm outside the toolkit (0.695 / 0.845 / 0.85 at steps 10 / 20 / 30). In a separate 512-prompt smoke run, dynamic sampling filtered 74-81% of groups per step, as expected for a low starting reward.

Note on dependencies

With uv sync, the lock file installs sagemaker-train 1.24.0 (no training_session module) and botocore 1.43.106 (no training-session operations), so neither train_grpo.py nor train_dapo.py runs. SETUP.md installs with uv pip install sagemaker-train, which should pick up current versions; I have not verified that path. I ran with sagemaker-train 1.28.0 and botocore 1.43.117. Not changed in this PR.

[Linbo-Liu]

is this a bit over-engineered?