diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 6684c03690..3aabf9a2ce 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -52,7 +52,69 @@ jobs: strategy: fail-fast: false matrix: - info: [{"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_short.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_sglang_config.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_sglang_config_distributed.py"}] + info: [{"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_ppo_critic_only_short.py"}] + defaults: + run: + working-directory: ${{ github.workspace }} + env: + GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} + WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} + SLIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} + SLIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} + SLIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} + SLIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + + - name: Install + shell: bash + run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages + + + - name: Execute + shell: bash + run: | + TEST_PATH="${{ matrix.info.test_file }}" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + if [ "${{ matrix.info.num_gpus }}" = "0" ]; then + python "$TEST_PATH" + else + python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" + fi + + e2e-test-sglang-config: + + if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-sglang-config')) + + + runs-on: self-hosted + container: + image: slimerl/slime:latest + options: > + --gpus all + --ipc=host + --shm-size=16g + --ulimit memlock=-1 + --ulimit stack=67108864 + --memory=0 + --memory-swap=0 + -e http_proxy=$http_proxy + -e https_proxy=$https_proxy + -e HTTP_PROXY=$HTTP_PROXY + -e HTTPS_PROXY=$HTTPS_PROXY + -v /mnt/nvme0n1/slime_ci:/data/slime_ci + -v /mnt/nvme0n1/slime_ci/models:/root/models + -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets + + strategy: + fail-fast: false + matrix: + info: [{"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_sglang_config.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_sglang_config_distributed.py"}, {"num_gpus": 8, "test_file": "test_sglang_config_mixed_offload.py"}, {"num_gpus": 8, "test_file": "test_sglang_config_mixed_offload_ft.py"}] defaults: run: working-directory: ${{ github.workspace }} diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 6ecc4376fb..666b322dd9 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -4,8 +4,16 @@ 'tests': [ {'test_file': 'test_qwen2.5_0.5B_gsm8k_async_short.py', 'num_gpus': 4}, {'test_file': 'test_qwen2.5_0.5B_gsm8k_short.py', 'num_gpus': 4}, + {'test_file': 'test_qwen2.5_0.5B_ppo_critic_only_short.py', 'num_gpus': 4}, + ], + }, + 'e2e-test-sglang-config': { + 'label': 'run-ci-sglang-config', + 'tests': [ {'test_file': 'test_qwen2.5_0.5B_sglang_config.py', 'num_gpus': 8}, {'test_file': 'test_qwen2.5_0.5B_sglang_config_distributed.py', 'num_gpus': 8}, + {'test_file': 'test_sglang_config_mixed_offload.py', 'num_gpus': 8}, + {'test_file': 'test_sglang_config_mixed_offload_ft.py', 'num_gpus': 8}, ], }, 'e2e-test-megatron': { diff --git a/slime/backends/megatron_utils/model.py b/slime/backends/megatron_utils/model.py index dd3724424d..4995b948b3 100644 --- a/slime/backends/megatron_utils/model.py +++ b/slime/backends/megatron_utils/model.py @@ -655,7 +655,7 @@ def train( # TODO: figure out why KL is not exactly zero when using PPO loss with KL clipping, and whether this is expected behavior or a bug. assert log_dict["train/ppo_kl"] < 1e-8, f"{log_dict=}" if accumulated_step_id == 0 and "train/kl_loss" in log_dict: - assert log_dict["train/kl_loss"] == 0.0, f"{log_dict=}" + assert log_dict["train/kl_loss"] < 1e-8, f"{log_dict=}" logger.info(f"{role_tag}step {accumulated_step_id}: {log_dict}") diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py index 66503bb6f2..bed0cf0b2c 100644 --- a/slime/backends/sglang_utils/sglang_engine.py +++ b/slime/backends/sglang_utils/sglang_engine.py @@ -90,21 +90,6 @@ def _wait_server_healthy(base_url, api_key, is_process_alive): time.sleep(2) - # use flush_cache to make sure the working queue is empty, so that we can do offload - while True: - try: - response = session.get(f"{base_url}/flush_cache", headers=headers) - if response.status_code == 200: - break - - except requests.RequestException: - pass - - if not is_process_alive(): - raise Exception("Server process terminated unexpectedly.") - - time.sleep(2) - class SGLangEngine(RayActor): def __init__( @@ -370,6 +355,17 @@ def resume_memory_occupation(self, tags: list[str] = None): def check_weights(self, action: str): return self._make_request("weights_checker", {"action": action}) + def update_weights_from_disk(self, model_path: str, load_format: str | None = None): + """Reload weights from *model_path* without restarting the engine. + + Used for non-updatable (frozen) models that overlap with megatron: + after offload, weights are restored from disk instead of CPU cache. + """ + payload = {"model_path": model_path} + if load_format is not None: + payload["load_format"] = load_format + return self._make_request("update_weights_from_disk", payload) + def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name, backend): return self._make_request( "init_weights_update_group", diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 5e41fbfd82..a7f3a191fe 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -53,6 +53,8 @@ class ServerGroup: rank_offset: int = 0 # cumulative engine count before this group gpu_offset: int = 0 # cumulative GPU count before this group sglang_overrides: dict = dataclasses.field(default_factory=dict) + needs_offload: bool = False # True when this group's GPUs overlap with megatron + model_path: str | None = None # checkpoint path for update_weights_from_disk router_ip: str | None = None router_port: int | None = None @@ -174,17 +176,35 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis def offload(self): """Fire release_memory_occupation on all engines (non-blocking). - Returns a list of Ray ObjectRefs. + Returns a list of Ray ObjectRefs. Skipped for groups that do not + overlap with megatron GPUs (``needs_offload=False``). """ + if not self.needs_offload: + return [] return [engine.release_memory_occupation.remote() for engine in self.engines if engine is not None] def onload(self, tags: list[str] | None = None): """Fire resume_memory_occupation on all engines (non-blocking). - Returns a list of Ray ObjectRefs. + Returns a list of Ray ObjectRefs. Skipped for groups that do not + overlap with megatron GPUs (``needs_offload=False``). """ + if not self.needs_offload: + return [] return [engine.resume_memory_occupation.remote(tags=tags) for engine in self.engines if engine is not None] + def onload_weights_from_disk(self): + """Reload weights from ``model_path`` for non-updatable groups. + + Used instead of ``resume_memory_occupation(tags=[WEIGHTS])`` so that + CPU memory is not consumed by offloaded weight copies. + """ + if not self.needs_offload or not self.model_path: + return [] + return [ + engine.update_weights_from_disk.remote(self.model_path) for engine in self.engines if engine is not None + ] + @dataclasses.dataclass class RolloutServer: @@ -261,20 +281,32 @@ def recover(self): # Post-recovery: offload then onload weights for newly created engines. release_handles = [] - new_engines_all = [] - for g, dead_indices in zip(self.server_groups, dead_per_group, strict=True): + updatable_new_engines = [] + non_updatable_groups_engines: list[tuple[str, list]] = [] + for g, dead_indices in zip(self.engine_groups, dead_per_group, strict=True): logger.info(f"Recovered {g.num_new_engines} dead rollout engines (worker_type={g.worker_type})") assert g.num_new_engines == len(dead_indices), "num_new_engines does not match dead_indices length" - if g.args.offload_rollout and dead_indices: + if g.needs_offload and dead_indices: new_engines = [g.all_engines[i] for i in dead_indices] release_handles.extend(engine.release_memory_occupation.remote() for engine in new_engines) - new_engines_all.extend(new_engines) + if self.update_weights: + updatable_new_engines.extend(new_engines) + elif g.model_path: + non_updatable_groups_engines.append((g.model_path, new_engines)) if release_handles: ray.get(release_handles) - ray.get( - [engine.resume_memory_occupation.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]) for engine in new_engines_all] - ) + # Resume GPU memory for all engines that need offload. + all_resume_engines = updatable_new_engines[:] + for _model_path, engines in non_updatable_groups_engines: + all_resume_engines.extend(engines) + if all_resume_engines: + ray.get( + [ + engine.resume_memory_occupation.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]) + for engine in all_resume_engines + ] + ) def offload(self): """Release memory occupation across all groups (concurrent).""" @@ -290,6 +322,28 @@ def onload(self, tags: list[str] | None = None): handles.extend(g.onload(tags)) return ray.get(handles) if handles else [] + def onload_weights(self): + """Restore weights for offloaded groups. + + All groups resume from CPU cache via ``resume_memory_occupation``. + For updatable servers, weights will be overwritten by + ``update_weights`` shortly after. For non-updatable servers the + CPU backup already contains the correct (unchanged) weights. + """ + handles = [] + for g in self.engine_groups: + if not g.needs_offload: + continue + handles.extend(g.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS])) + return ray.get(handles) if handles else [] + + def onload_kv(self): + """Resume KV cache and CUDA graphs for offloaded groups.""" + handles = [] + for g in self.engine_groups: + handles.extend(g.onload(tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH])) + return ray.get(handles) if handles else [] + @ray.remote class RolloutManager: @@ -444,10 +498,12 @@ def onload(self, tags: list[str] | None = None): srv.onload(tags) def onload_weights(self): - self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS]) + for srv in self.servers.values(): + srv.onload_weights() def onload_kv(self): - self.onload(tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH]) + for srv in self.servers.values(): + srv.onload_kv() def recover_updatable_engines(self): """Restart any dead rollout engines and update num_new_engines for update_weights detection. @@ -887,6 +943,30 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool return router_ip, router_port +def _compute_rollout_offset(args) -> int: + """Offset (in PG bundle slots) where rollout GPUs start.""" + if args.debug_train_only or args.debug_rollout_only or args.colocate: + return 0 + if args.critic_train_only: + return args.critic_num_nodes * args.critic_num_gpus_per_node + offset = args.actor_num_nodes * args.actor_num_gpus_per_node + if args.use_critic: + offset += args.critic_num_nodes * args.critic_num_gpus_per_node + return offset + + +def _compute_megatron_num_gpus(args) -> int: + """Total number of megatron (actor + critic) GPU slots in the placement group.""" + if args.debug_rollout_only: + return 0 + if args.critic_train_only: + return args.critic_num_nodes * args.critic_num_gpus_per_node + num = args.actor_num_nodes * args.actor_num_gpus_per_node + if args.use_critic: + num += args.critic_num_nodes * args.critic_num_gpus_per_node + return num + + def start_rollout_servers(args, pg) -> dict[str, RolloutServer]: """Start rollout servers: one per model, each with its own router. @@ -906,6 +986,10 @@ def start_rollout_servers(args, pg) -> dict[str, RolloutServer]: gpu_offset = 0 engine_offset = 0 + # Compute megatron GPU range for per-group offload decisions. + rollout_pg_offset = _compute_rollout_offset(args) + megatron_num_gpus = _compute_megatron_num_gpus(args) + for model_idx, model_cfg in enumerate(config.models): model_cfg.resolve(args) @@ -926,6 +1010,17 @@ def start_rollout_servers(args, pg) -> dict[str, RolloutServer]: num_gpu_per_engine_local = min(gpus_per_engine, args.num_gpus_per_node) num_engines = group_cfg.num_gpus // num_gpu_per_engine_local + # Only offload groups whose GPUs overlap with megatron. + group_abs_start = rollout_pg_offset + gpu_offset + needs_offload = args.offload_rollout and group_abs_start < megatron_num_gpus + overrides = dict(group_cfg.overrides) + if args.offload_rollout and not needs_offload: + overrides.setdefault("enable_memory_saver", False) + logger.info( + f"Engine group '{group_cfg.worker_type}' gpu_offset={gpu_offset} " + f"(abs={group_abs_start}): needs_offload={needs_offload}" + ) + group = ServerGroup( args=args, pg=pg, @@ -935,7 +1030,9 @@ def start_rollout_servers(args, pg) -> dict[str, RolloutServer]: worker_type=group_cfg.worker_type, rank_offset=engine_offset, gpu_offset=gpu_offset, - sglang_overrides=group_cfg.overrides, + sglang_overrides=overrides, + needs_offload=needs_offload, + model_path=overrides.get("model_path", args.hf_checkpoint), router_ip=router_ip, router_port=router_port, ) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 4ce9d3250d..64a6d61902 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -1521,6 +1521,15 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: def slime_validate_args(args): args.eval_datasets = _resolve_eval_datasets(args) + if args.critic_train_only: + if not args.use_critic: + raise ValueError("--critic-train-only requires --use-critic (or --advantage-estimator ppo).") + if args.actor_num_nodes != 0 or args.actor_num_gpus_per_node != 0: + raise ValueError( + "--critic-train-only requires --actor-num-nodes 0 --actor-num-gpus-per-node 0, " + f"but got actor_num_nodes={args.actor_num_nodes}, actor_num_gpus_per_node={args.actor_num_gpus_per_node}." + ) + if args.kl_coef != 0 or args.use_kl_loss: if not os.path.exists(args.ref_load): raise FileNotFoundError(f"ref_load {args.ref_load} does not exist, please check the path.") diff --git a/tests/test_sglang_config_mixed_offload.py b/tests/test_sglang_config_mixed_offload.py new file mode 100644 index 0000000000..7451331359 --- /dev/null +++ b/tests/test_sglang_config_mixed_offload.py @@ -0,0 +1,158 @@ +"""E2E test: mixed offload with updatable + frozen models. + +Deploys two models via --sglang-config in colocate mode: + - "actor": update_weights=true, 4 GPUs → overlaps with megatron, gets offloaded + and weights updated from training. + - "ref": update_weights=false, 4 GPUs → overlaps with megatron, gets offloaded + and weights restored from disk (update_weights_from_disk). + +Key coverage: + - Per-group needs_offload (both overlap with megatron in colocate mode) + - update_weights_from_disk for frozen model + - Selective flush_cache (only for offloaded / updatable engines) + - Offload/onload cycle completes without crash +""" + +import os +import tempfile + +import slime.utils.external_utils.command_utils as U + +TIGHT_DEVICE_MEMORY = U.get_bool_env_var("SLIME_TEST_TIGHT_DEVICE_MEMORY", "1") + +MODEL_NAME = "Qwen2.5-0.5B-Instruct" +MODEL_TYPE = "qwen2.5-0.5B" +NUM_GPUS = 8 + +# Two models on 8 GPUs (colocate): actor gets weight updates, ref is frozen. +SGLANG_CONFIG_YAML = """\ +sglang: + - name: actor + update_weights: true + engine_groups: + - worker_type: regular + num_gpus: 4 + num_gpus_per_engine: 1 + - name: ref + update_weights: false + engine_groups: + - worker_type: regular + num_gpus: 4 + num_gpus_per_engine: 1 +""" + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/gsm8k") + + +def execute(): + config_file = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", prefix="sglang_mixed_offload_", delete=False) + config_file.write(SGLANG_CONFIG_YAML) + config_file.flush() + config_path = config_file.name + + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " + + rollout_args = ( + "--prompt-data /root/datasets/gsm8k/train.parquet " + "--input-key messages " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + "--num-rollout 3 " + "--rollout-batch-size 8 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 512 " + "--rollout-temperature 0.8 " + "--global-batch-size 32 " + ) + + eval_args = ( + "--eval-interval 20 " + "--eval-prompt-data gsm8k /root/datasets/gsm8k/test.parquet " + "--n-samples-per-eval-prompt 1 " + "--eval-max-response-len 512 " + "--eval-top-k 1 " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 4096 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + sglang_args = ( + "--rollout-num-gpus-per-engine 1 " + f"--sglang-mem-fraction-static {0.6 if TIGHT_DEVICE_MEMORY else 0.7} " + "--sglang-cuda-graph-max-bs 32 " + f"--sglang-config {config_path} " + ) + + ci_args = "--ci-test " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 8 " + "--colocate " + "--megatron-to-hf-mode bridge " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{eval_args} " + f"{sglang_args} " + f"{ci_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + +if __name__ == "__main__": + prepare() + os.environ.pop("http_proxy", None) + os.environ.pop("https_proxy", None) + os.environ.pop("HTTP_PROXY", None) + os.environ.pop("HTTPS_PROXY", None) + execute() diff --git a/tests/test_sglang_config_mixed_offload_ft.py b/tests/test_sglang_config_mixed_offload_ft.py new file mode 100644 index 0000000000..48b9063899 --- /dev/null +++ b/tests/test_sglang_config_mixed_offload_ft.py @@ -0,0 +1,165 @@ +"""E2E test: mixed offload with fault tolerance. + +Same two-model layout as test_sglang_config_mixed_offload.py but with +fault tolerance enabled. --ci-test triggers a simulated engine crash +on the updatable (actor) server, testing: + - Health monitor detects crash and marks engine as None + - RolloutServer.recover() restarts the dead engine + - Updatable engines: offload → resume_memory_occupation → update_weights + - Non-updatable engines: offload → update_weights_from_disk + - Training continues after recovery +""" + +import os +import tempfile + +import slime.utils.external_utils.command_utils as U + +TIGHT_DEVICE_MEMORY = U.get_bool_env_var("SLIME_TEST_TIGHT_DEVICE_MEMORY", "1") + +MODEL_NAME = "Qwen2.5-0.5B-Instruct" +MODEL_TYPE = "qwen2.5-0.5B" +NUM_GPUS = 8 + +# Two models on 8 GPUs (colocate): actor gets weight updates, ref is frozen. +SGLANG_CONFIG_YAML = """\ +sglang: + - name: actor + update_weights: true + engine_groups: + - worker_type: regular + num_gpus: 4 + num_gpus_per_engine: 1 + - name: ref + update_weights: false + engine_groups: + - worker_type: regular + num_gpus: 4 + num_gpus_per_engine: 1 +""" + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/gsm8k") + + +def execute(): + config_file = tempfile.NamedTemporaryFile( + mode="w", suffix=".yaml", prefix="sglang_mixed_offload_ft_", delete=False + ) + config_file.write(SGLANG_CONFIG_YAML) + config_file.flush() + config_path = config_file.name + + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " + + rollout_args = ( + "--prompt-data /root/datasets/gsm8k/train.parquet " + "--input-key messages " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + "--num-rollout 3 " + "--rollout-batch-size 8 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 512 " + "--rollout-temperature 0.8 " + "--global-batch-size 32 " + ) + + eval_args = ( + "--eval-interval 20 " + "--eval-prompt-data gsm8k /root/datasets/gsm8k/test.parquet " + "--n-samples-per-eval-prompt 1 " + "--eval-max-response-len 512 " + "--eval-top-k 1 " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 4096 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + sglang_args = ( + "--rollout-num-gpus-per-engine 1 " + f"--sglang-mem-fraction-static {0.6 if TIGHT_DEVICE_MEMORY else 0.7} " + "--sglang-cuda-graph-max-bs 32 " + f"--sglang-config {config_path} " + ) + + ci_args = "--ci-test " + + fault_tolerance_args = ( + "--use-fault-tolerance " + "--rollout-health-check-interval 5 " + "--rollout-health-check-timeout 10 " + "--rollout-health-check-first-wait 0 " + ) + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 8 " + "--colocate " + "--megatron-to-hf-mode bridge " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{eval_args} " + f"{sglang_args} " + f"{ci_args} " + f"{fault_tolerance_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + +if __name__ == "__main__": + prepare() + os.environ.pop("http_proxy", None) + os.environ.pop("https_proxy", None) + os.environ.pop("HTTP_PROXY", None) + os.environ.pop("HTTPS_PROXY", None) + execute() diff --git a/train.py b/train.py index 7e7f752b3a..15b3ca8b24 100644 --- a/train.py +++ b/train.py @@ -23,10 +23,11 @@ def train(args): ray.get(rollout_manager.onload_weights.remote()) # always update weight first so that sglang has the loaded weights from training. - actor_model.update_weights() + if not args.critic_train_only: + actor_model.update_weights() - if args.check_weight_update_equal: - ray.get(rollout_manager.check_weights.remote(action="compare")) + if args.check_weight_update_equal: + ray.get(rollout_manager.check_weights.remote(action="compare")) if args.offload_rollout: ray.get(rollout_manager.onload_kv.remote()) @@ -39,15 +40,18 @@ def offload_train(rollout_id): if args.offload_train: if args.use_critic: critic_model.offload() - if rollout_id >= args.num_critic_only_steps: + if rollout_id >= args.num_critic_only_steps and not args.critic_train_only: actor_model.offload() else: actor_model.offload() else: - actor_model.clear_memory() + if args.critic_train_only: + critic_model.clear_memory() + else: + actor_model.clear_memory() def save(rollout_id): - if (not args.use_critic) or (rollout_id >= args.num_critic_only_steps): + if (not args.use_critic) or (rollout_id >= args.num_critic_only_steps and not args.critic_train_only): actor_model.save_model( rollout_id, force_sync=rollout_id == args.num_rollout - 1, @@ -73,7 +77,7 @@ def save(rollout_id): if args.use_critic: critic_train_handle = critic_model.async_train(rollout_id, rollout_data_ref) - if rollout_id >= args.num_critic_only_steps: + if rollout_id >= args.num_critic_only_steps and not args.critic_train_only: ray.get(actor_model.async_train(rollout_id, rollout_data_ref)) ray.get(critic_train_handle) else: @@ -85,7 +89,8 @@ def save(rollout_id): offload_train(rollout_id) if args.offload_rollout: ray.get(rollout_manager.onload_weights.remote()) - actor_model.update_weights() + if not args.critic_train_only: + actor_model.update_weights() if args.offload_rollout: ray.get(rollout_manager.onload_kv.remote()) diff --git a/train_async.py b/train_async.py index fc6c17d58c..182309cb08 100644 --- a/train_async.py +++ b/train_async.py @@ -22,10 +22,11 @@ def train(args): actor_model, critic_model = create_training_models(args, pgs, rollout_manager) # always update weight first so that sglang has the loaded weights from training. - actor_model.update_weights() + if not args.critic_train_only: + actor_model.update_weights() - if args.check_weight_update_equal: - ray.get(rollout_manager.check_weights.remote(action="compare")) + if args.check_weight_update_equal: + ray.get(rollout_manager.check_weights.remote(action="compare")) # async train loop. rollout_data_next_future = rollout_manager.generate.remote(args.start_rollout_id) @@ -40,17 +41,18 @@ def train(args): if args.use_critic: critic_train_handle = critic_model.async_train(rollout_id, rollout_data_curr_ref) - if rollout_id >= args.num_critic_only_steps: + if rollout_id >= args.num_critic_only_steps and not args.critic_train_only: ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref)) ray.get(critic_train_handle) else: ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref)) if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout): - actor_model.save_model( - rollout_id, - force_sync=rollout_id == args.num_rollout - 1, - ) + if not args.critic_train_only: + actor_model.save_model( + rollout_id, + force_sync=rollout_id == args.num_rollout - 1, + ) if args.use_critic: critic_model.save_model( rollout_id, @@ -63,7 +65,8 @@ def train(args): # sync generate before update weights to prevent update weight in the middle of generation rollout_data_curr_ref = ray.get(x) if (x := rollout_data_next_future) is not None else None rollout_data_next_future = None - actor_model.update_weights() + if not args.critic_train_only: + actor_model.update_weights() if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch): ray.get(rollout_manager.eval.remote(rollout_id))