Feature/cispo algorithm#1
Merged
Merged
Conversation
…parallelism (areal-project#841) Prepare archon infrastructure for pipeline parallelism (PP) support: - Add pp dimension to ArchonParallelDims with pp_enabled property - Add dp_replicate placeholder for future HSDP support - Replace reshard_after_forward bool with reshard_after_forward_policy string ("default"/"always"/"never") to handle PP-specific resharding - Fix MoE init_weights to require buffer_device (remove optional None) - Add shard_placement_fn for MoE experts when dp_mod_ep * ep > num_experts - Improve comments for FSDP resharding optimization on final layers
…l-project#844) Split module.py into separate files to allow FSDP engine to import tree attention functionality without requiring Megatron dependencies: - constants.py: shared BLOCK_SIZE constant (lightweight, no deps) - module_fsdp.py: FSDP-only functionality (PyTorch only) - module_megatron.py: Megatron-specific functionality (mbridge, megatron.core) - module.py: backward-compatible re-exports with conditional Megatron imports
…real-project#840) * feat(workflow): add Anthropic Messages API support for RL training Enable Anthropic-compatible agents to run through AReaL's proxy server for reinforcement learning workflows. This allows using Claude-style agents with the existing RL infrastructure. Key changes: - Add /v1/messages endpoint to proxy server with LiteLLM-based conversion - Introduce MathAgent workflow using Anthropic Messages API format - Add anthropic and litellm dependencies for API conversion --------- Co-authored-by: 博惟 <bowei.fw@antgroup.com>
* fix concat logic when seq end with pad_token * remove unittest skip --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This pull request adds support for dynamically loading a reward function from a string path in `VisionRLVRWorkflow`, making it consistent with the parent `RLVRWorkflow`. The change is functionally correct and fixes a bug where using a string path for `reward_fn` would previously lead to an `AttributeError`.
This pull request correctly fixes a bug in the `wandb` initialization logic within `areal/utils/stats_logger.py`. Previously, the environment variables `WANDB_API_KEY` and `WANDB_BASE_URL` were assigned each other's values from the configuration. The submitted change rectifies this by correctly mapping the configuration values to their respective environment variables.
…upport (areal-project#849) Replace manual safetensors handling with PyTorch DCP infrastructure for HuggingFace format checkpoints. Previously, save/load required gathering full state dict to rank 0, causing OOM on large MoE models. Now uses HuggingFaceStorageWriter/Reader for distributed operations. Key changes: - Add MoEStateDictAdapter base class with DTensor-aware expert weight conversion (3D grouped <-> 2D individual), supporting EP and ETP placements via local expert index calculation - Fix in-place tensor ops (masked_fill_, scatter_) for torch.compile and checkpointing compatibility - Large MoE checkpoint save/load without OOM (no full tensor gather) - Correct EP/ETP checkpoint handling with distributed expert weights - Incremental weight sync to inference engines via convert_single_to_hf() - Foundation for PP support: layer-indexed key mapping and extensible DTensor placement logic for per-stage conversion
…areal-project#843) * ci(docker): modernize dependency management and Docker build pipeline This PR overhauls the dependency management and Docker build infrastructure for improved reproducibility, faster builds, and better developer experience. Key changes: - **Dependency Management**: Migrate to uv-based workflow with locked dependencies (uv.lock). Reorganize pyproject.toml with platform-specific extras (cuda, sglang, vllm, megatron) and dev dependencies. This ensures consistent environments across development, CI, and production. - **Docker Build Optimization**: Restructure Dockerfile with layered caching strategy - install torch and C++ packages (apex, grouped_gemm, transformer_engine) BEFORE syncing pyproject.toml dependencies. This prevents expensive recompilation when only Python dependencies change. - **CI/CD Improvements**: - Add runner-heartbeat workflow to prevent GitHub self-hosted runner expiration (7-day schedule) - Fix build-docker-image workflow to preserve running builder instances - Update tag-release-image to build fresh images instead of re-tagging - Update test-areal to use venv-based Python environment - **Code Modernization**: - Replace TypeVar with PEP 695 type parameter syntax (Python 3.12+) - Replace deprecated asyncio.TimeoutError with TimeoutError - Rename areal/tools/profile to profiling_utils to avoid import conflicts - Update vLLM imports for v0.14.0 compatibility - **Documentation**: Rewrite installation guide for uv-based workflow with clear instructions for CUDA/non-CUDA environments. * fix flash-attn version check
…t#858) * test: add unit tests for abort loop condition fix in RemoteInfEngine Add comprehensive tests to verify the fix for the bug where the generation loop would prematurely exit when stop_reason='abort'. The bug was caused by using gconfig.max_new_tokens (which gets decremented in the loop) instead of the original value in the loop condition. Tests include: - test_abort_should_continue_loop_with_fix: verify fix works correctly - test_abort_prematurely_exits_without_fix: reproduce the original bug - test_abort_multiple_times_continues_with_fix: multiple abort scenarios - test_abort_reaches_length_limit_with_fix: length limit handling - Additional tests for tool_calls, length stop reasons, and gconfig behavior Reference: commit 445dca7216e66c24d7a5b47a272761ea81749f8d Co-Authored-By: Claude (claude-opus-4-5) <noreply@anthropic.com> * refactor: use real RemoteInfEngine with mocked arequest_with_retry Rewrite tests to directly use RemoteInfEngine instance instead of simulating the loop logic. This provides more realistic testing by: - Creating actual RemoteInfEngine with MockBackend - Mocking arequest_with_retry to avoid needing a real SGLang server - Testing the actual agenerate() method behavior The tests now verify that the fix works correctly in the real codebase, not just in isolated simulation. Co-Authored-By: Claude (claude-opus-4-5) <noreply@anthropic.com> --------- Co-authored-by: Claude (claude-opus-4-5) <noreply@anthropic.com>
…l-project#859) Dynamically compute thread control environment variables (OMP_NUM_THREADS, MKL_NUM_THREADS, etc.) based on allocated CPU cores to prevent thread explosion in multi-process training scenarios. Changes: - Add get_thread_env_vars() function in launcher utils - Update SchedulingSpec.cpu default from 4 to 8 - Apply thread env vars in Local, Ray, and Slurm launchers - Support user override via SchedulingSpec.env_vars or os.environ
…real-project#860) Optimize Archon engine initialization to use meta device + FSDP + DCP loading pattern, reducing peak memory from N × model_size to ~1/N × model_size per rank. Changes: - Create model structure on meta device (no memory allocation) - Apply FSDP/TP/EP parallelization to meta device model - Materialize tensors with to_empty() after FSDP - Load weights via DCP (each rank reads only its shard) - Add init_buffers() to BaseArchonModel for buffer-only initialization - Separate init_weights() (parameters) from init_buffers() (rope_cache)
…zation (areal-project#862) When training large models, loading weights directly onto each GPU can cause OOM during initialization. This adds `fsdp.memory_efficient_load` config option that initializes model structure on CPU, then rank 0 loads pretrained weights and broadcasts to all ranks after FSDP sharding. This reduces peak GPU memory during model loading at the cost of slightly longer initialization time.
…ect#861) * ci(docker): add DeepSeek-V3 kernels and modernize tooling Enhance Docker image with DeepSeek-V3 optimization packages (FlashMLA, DeepGEMM, DeepEP, flash-linear-attention) to enable MoE training with optimal performance on Hopper GPUs. Also modernize development tooling and resolve dependency conflicts. Dockerfile changes: - Add FlashMLA, DeepGEMM, DeepEP, flash-linear-attention for DeepSeek-V3 - Replace nvm with fnm (Fast Node Manager) for faster Node.js setup - Move Node.js/npm tooling to STAGE 2.5 after CUDA compilation - Add AI coding tools: opencode, @openai/codex, @google/gemini-cli Dependency changes: - Add torchao>=0.15.0 for PyTorch quantization and optimization - Add claude-agent-sdk for Claude Agent SDK integration - Replace litellm + litellm-proxy-extras with litellm[proxy]>=1.81.3 to resolve uvicorn version conflict - Add openai and soundfile to override-dependencies Validation and documentation: - Add DeepSeek-V3 package tests to validate_docker_installation.py - Document GPU requirements (SM90+ for FlashMLA/DeepGEMM, SM80+ for DeepEP) * remove nvshm validation * fix fla version validation --------- Co-authored-by: 博惟 <bowei.fw@antgroup.com>
Set up comprehensive Claude Code integration with domain-specific agents, commands, rules, and skills for the codebase. Key changes: - Add 7 specialized agents (algorithm, FSDP, Megatron, Archon experts) - Add /pr-review and /gen-commit-msg commands - Add rules for code style, API config, distributed, and testing - Add skills for adding datasets, rewards, workflows, and debugging - Update CLAUDE.md with project overview and guidelines - Configure settings.json with project-specific permissions
…areal-project#864) Implement Pipeline Parallelism for the Archon training engine using PyTorch's PipelineStage and Schedule1F1B, enabling model sharding across multiple GPU stages for training larger models. - Add pipeline infrastructure for balanced layer distribution and model splitting - Extend engine to manage multiple model parts with pipelined forward/backward - Extend device mesh to support PP dimension (world_size = pp × dp × cp × tp) - Update DCP and HuggingFace checkpoint save/load for pipeline-sharded models - Adapt Qwen2/Qwen3 models to handle PP tensor shapes and partial modules - Fix attention mask caching incompatibility with PP microbatches
* feat(utils): add dynamic PYTHONPATH in get_env_vars Dynamically build PYTHONPATH from current sys.path to ensure spawned worker processes can import modules from the same locations as the parent process. This is essential for deserializing custom config classes defined outside the areal package (e.g., in examples/). The dynamic PYTHONPATH includes: - Current working directory (first) - All sys.path entries (excluding .zip files) - Original PYTHONPATH environment variable Co-Authored-By: Claude (claude-opus-4-5) <noreply@anthropic.com> * Update areal/utils/launcher.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Remove redundant --------- Co-authored-by: Claude (claude-opus-4-5) <noreply@anthropic.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Wei Fu <36355462+garrett4wade@users.noreply.github.com>
…eal-project#873) Comprehensively restructure Claude agent configuration with expanded content and improved architecture. Add launcher-scheduler-expert agent and add-unit-tests skill to establish complete agent ecosystem. Key changes: - Engine expert refactoring: Complete rewrite and expansion of engine agents (archon/fsdp/megatron) with 400+ additional lines of detailed guidance - New launcher-scheduler-expert: 185-line agent for cluster launching and resource scheduling (Slurm/Ray/Kubernetes) - New add-unit-tests skill: 228-line guided workflow for test development - Updated CLAUDE.md: Complete agent table with stage-by-stage guidance - Symbol cleanup: Replace all non-ASCII symbols (→, –, ✓, Greek letters) with ASCII equivalents across agent documentation - Architecture alignment: Rename engine experts from generic "expert" to specific "engine-expert" for clarity and consistency
…real-project#872) * lazy attn mask for tree training * Raise error for padded_size.
…real-project#875) Consolidate scattered infrastructure modules into a unified `areal/infra` package for better organization and discoverability. Key changes: - Rename `areal/core` → `areal/infra` to better reflect infrastructure role - Move `areal/controller` → `areal/infra/controller` for logical grouping - Move `areal/platforms` → `areal/infra/platforms` to colocate platform abstractions - Move vLLM extensions from `areal/thirdparty/vllm` → `areal/engine/vllm_ext` to align with other engine extensions - Add `.claude/commands/create-pr.md` for automated PR creation workflow - Add launcher-scheduler expert agent documentation - Update all imports across 135 files to reflect new structure
…_bash_cmds` arguments. (areal-project#876)
…eal-project#877) Enable pipeline parallelism (PP > 1) to work correctly with RL training by fixing data broadcast and weight synchronization issues. Key changes: - Replace cp_tp mesh with pp_cp_tp mesh to include PP in context_and_model_parallel_group - Add per-PP-stage weight update groups for XCCL weight sync - Add DistributedLock to serialize NCCL group initialization - Add PP + weight_tying validation (unsupported combination) - Fix MegatronEngine missing engine_lock in init
This PR implements a Triton kernel that supports tree attention.
areal-project#882) Extract forward/backward execution and weight synchronization logic from ArchonEngine into dedicated modules for better separation of concerns. Key changes: - Add ForwardBackwardRunner abstraction with SequentialRunner and PipelinedRunner implementations - Extract weight sync logic to WeightSyncState and dedicated functions - Simplify planner agent guidelines with clearer structure
… submodules (areal-project#1221) Break monolithic data_proxy/backend.py into properly-separated submodules with a protocol-based backend abstraction, preparing the inference_service package structure for weight update integration. Key changes: - Add InfBridgeBackend protocol in inference_service/backend.py - Extract SGLangBridgeBackend into sglang/bridge.py - Extract VLLMBridgeBackend into vllm/bridge.py - Move inf_bridge.py from data_proxy/ to inference_service/ - Create sglang/ and vllm/ as proper Python packages - Update all consumer imports (app.py, 9 test files)
* feat: add awex backend for weight update Add weight update infrastructure using awex for cross-engine weight synchronization via NCCL P2P between FSDP training workers and SGLang inference servers. Currently supports FSDP (pure DP) + SGLang (pure DP) only. Key changes: - Add weight update gateway with HTTP API for connect/update/disconnect lifecycle - Add FSDP and SGLang awex adapters for parameter shard mapping and P2P transfer - Add NCCL group management for cross-process-group communication - Add weight update controller for client-side orchestration - Refactor SGLang inference service with custom launch_server and scheduler - Extend RPC serialization to handle torch.dtype and Enum types - Add awex==0.5.0 dependency - Add comprehensive unit and integration tests * fix: address PR review feedback - Fail fast with 500 on non-dict parallelism result instead of returning empty dict (awex.py) - Log warnings on shutdown exceptions instead of silently passing (controller.py) - Warn on nameless parameter metadata in merge (gateway/app.py) * chore: skip ray serialization test * chore: fix tests * chore: fix broken inf service path
…g loop (areal-project#1163) Eliminate unnecessary GPU↔CPU parameter residency churn when offloading is enabled, addressing three sources of redundant transitions: 1. Critic double cycle: keep critic onloaded across compute_values and ppo_update instead of two separate onload/offload cycles. 2. Actor paused window: keep actor onloaded from ppo_update through update_weights + save + checkpoint, offload once before eval. 3. Reentrant _offload_aware_context: nested calls only perform actual onload/offload at the outermost level, preventing redundant transitions when engine methods (save, export_stats) are called while already onloaded. Verified on 4×H100 with FSDP2 + torch_memory_saver offload/onload cycles. Closes areal-project#1158
…g CP logits (areal-project#1223) Co-authored-by: 博惟 <bowei.fw@antgroup.com>
* feat: add disk-mode weight update flow to gateway Support disk-based weight transfer for weight-update pairs, including save-to-disk on train workers and load/update on inference workers. Ensure inference pause/continue is wrapped in a context manager so generation resumes even when updates fail. Key changes: - add disk mode fields to connect request and pair metadata - implement disk transfer path with save/update_weights_from_disk and LoRA loading - add controller and gateway tests for disk mode behavior * fix: validate disk-mode gateway connect inputs Fail fast on disk-mode misconfiguration so weight updates do not silently write to local relative paths instead of shared storage. Key changes: - reject empty disk-mode save_path values - require absolute save_path values for disk mode - require lora_name when disk-mode LoRA is enabled - add gateway tests for rejected connect requests --------- Co-authored-by: Wentai Zhang <rchardx@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
19 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Implementation of CISPO Algorithm with experiment test scripts