Closed
Add comprehensive PlantUML system architecture diagram for AReaL#2
Conversation
Copilot
AI
changed the title
[WIP] Add PlantUML high-level architecture diagram for AReaL system
Add comprehensive PlantUML system architecture diagram for AReaL
Jul 29, 2025
zhshgmail
pushed a commit
that referenced
this pull request
Oct 8, 2025
## Bugs Fixed ### Bug #1: Incorrect proximal_logprobs_t Length (CRITICAL) - File: areal/engine/sglang_remote.py - Issue: Length mismatch causing crashes during abort-resume - Fix: Changed EXTEND to REPLACE logic for previous tokens - Impact: Prevents crashes like "len(proximal_t)=7 vs len(output)=3" ### Bug #2: Misaligned Tensor Roll (CRITICAL) - File: areal/engine/ppo/actor.py - Issue: Off-by-one misalignment in loss computation - Fix: Added roll operation for old_logp to align with proximal_logprobs_t - Impact: Ensures correct importance weight calculation ## Changes - Modified: areal/engine/sglang_remote.py (lines 183, 240-282) - Modified: areal/engine/ppo/actor.py (lines 322-327) - Modified: areal/api/workflow_api.py (added recompute_all_proximal_t method) - Added: test_proximal_t_generation_simple.py (5 unit tests for Bug #1) - Added: test_recompute_timing.py (11 tests for recompute logic) - Added: CHANGES_SUMMARY.md (quick reference) - Added: docs/PROXIMAL_T_BUGS_COMPLETE_ANALYSIS.md (comprehensive documentation) ## Test Results All 42 tests pass: - test_proximal_t_generation_simple.py: 5/5 ✓ - test_recompute_timing.py: 11/11 ✓ - test_workflow_executor_wait.py: 26/26 ✓ 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
zhshgmail
pushed a commit
that referenced
this pull request
Oct 8, 2025
This commit fixes critical bugs in segment-wise decoupled PPO and adds symmetric importance weight clipping capability. ## Bug Fixes ### Bug #1: Incorrect proximal_logprobs_t Length (CRITICAL) - **File**: areal/engine/sglang_remote.py (lines 183, 240-282) - **Issue**: Length mismatch during abort-resume causing crashes - **Root cause**: Used EXTEND instead of REPLACE, duplicating values - **Fix**: Changed to REPLACE logic for previous tokens - **Impact**: Prevents crashes like "len(proximal_t)=7 vs len(output)=3" ### Bug #2: Incorrect old_logp Rolling (CRITICAL) - **File**: areal/engine/ppo/actor.py - **Issue**: old_logp was double-rolled causing misalignment - **Root cause**: old_logp already rolled in compute_logp() at line 93 - **Fix**: Removed incorrect second roll in grpo_loss_fn() - **Impact**: Ensures correct alignment for loss computation ## New Features ### behav_imp_weight_floor Parameter - **Motivation**: Original behav_imp_weight_cap only bounds high ratios, but stable importance sampling needs symmetric bounds - **Implementation**: Optional floor parameter for lower bound clipping - **Backward compatible**: floor=None maintains original behavior - **Usage**: - Symmetric: cap=10.0, floor=0.1 creates trust region [0.1, 10.0] - Asymmetric: Custom cap/floor ratios supported ### Files Modified **Config parameters:** - areal/api/cli_args.py: Added behav_imp_weight_floor - realhf/api/cli_args.py: Added behav_imp_weight_floor **Loss functions:** - areal/utils/functional.py: Updated ppo_actor_loss_fn - areal/engine/ppo/actor.py: Updated grpo_loss_fn (bug fix + floor) - realhf/impl/model/utils/ppo_functional.py: Updated actor_loss_fn - realhf/impl/model/interface/ppo_interface.py: Updated PPOInterface **Workflow infrastructure:** - areal/api/workflow_api.py: Added recompute_all_proximal_t() method ## Tests ### New Test Files - **test_proximal_t_generation_simple.py**: 5 tests for Bug #1 - **test_recompute_timing.py**: 11 tests for recompute logic - **test_behav_imp_weight_floor.py**: 10 tests for floor parameter - **test_grpo_loss_logp_alignment.py**: 4 tests guarding against Bug #2 ### Test Results All 56 tests pass: - test_workflow_executor_wait.py: 26/26 ✓ - test_proximal_t_generation_simple.py: 5/5 ✓ - test_recompute_timing.py: 11/11 ✓ - test_behav_imp_weight_floor.py: 10/10 ✓ - test_grpo_loss_logp_alignment.py: 4/4 ✓ ## Documentation - CHANGES_SUMMARY.md: Quick reference for all changes - docs/PROXIMAL_T_BUGS_COMPLETE_ANALYSIS.md: Comprehensive analysis 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
zhshgmail
pushed a commit
that referenced
this pull request
Oct 9, 2025
This commit fixes critical bugs in segment-wise decoupled PPO and adds symmetric importance weight clipping capability. ## Bug Fixes ### Bug #1: Incorrect proximal_logprobs_t Length (CRITICAL) - **File**: areal/engine/sglang_remote.py (lines 183, 240-282) - **Issue**: Length mismatch during abort-resume causing crashes - **Root cause**: Used EXTEND instead of REPLACE, duplicating values - **Fix**: Changed to REPLACE logic for previous tokens - **Impact**: Prevents crashes like "len(proximal_t)=7 vs len(output)=3" ### Bug #2: Incorrect old_logp Rolling (CRITICAL) - **File**: areal/engine/ppo/actor.py - **Issue**: old_logp was double-rolled causing misalignment - **Root cause**: old_logp already rolled in compute_logp() at line 93 - **Fix**: Removed incorrect second roll in grpo_loss_fn() - **Impact**: Ensures correct alignment for loss computation ## New Features ### behav_imp_weight_floor Parameter - **Motivation**: Original behav_imp_weight_cap only bounds high ratios, but stable importance sampling needs symmetric bounds - **Implementation**: Optional floor parameter for lower bound clipping - **Backward compatible**: floor=None maintains original behavior - **Usage**: - Symmetric: cap=10.0, floor=0.1 creates trust region [0.1, 10.0] - Asymmetric: Custom cap/floor ratios supported ### Files Modified **Config parameters:** - areal/api/cli_args.py: Added behav_imp_weight_floor - realhf/api/cli_args.py: Added behav_imp_weight_floor **Loss functions:** - areal/utils/functional.py: Updated ppo_actor_loss_fn - areal/engine/ppo/actor.py: Updated grpo_loss_fn (bug fix + floor) - realhf/impl/model/utils/ppo_functional.py: Updated actor_loss_fn - realhf/impl/model/interface/ppo_interface.py: Updated PPOInterface **Workflow infrastructure:** - areal/api/workflow_api.py: Added recompute_all_proximal_t() method ## Tests ### New Test Files - **test_proximal_t_generation_simple.py**: 5 tests for Bug #1 - **test_recompute_timing.py**: 11 tests for recompute logic - **test_behav_imp_weight_floor.py**: 10 tests for floor parameter - **test_grpo_loss_logp_alignment.py**: 4 tests guarding against Bug #2 ### Test Results All 56 tests pass: - test_workflow_executor_wait.py: 26/26 ✓ - test_proximal_t_generation_simple.py: 5/5 ✓ - test_recompute_timing.py: 11/11 ✓ - test_behav_imp_weight_floor.py: 10/10 ✓ - test_grpo_loss_logp_alignment.py: 4/4 ✓ ## Documentation - CHANGES_SUMMARY.md: Quick reference for all changes - docs/PROXIMAL_T_BUGS_COMPLETE_ANALYSIS.md: Comprehensive analysis 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
zhshgmail
pushed a commit
that referenced
this pull request
Oct 10, 2025
Implement segment-wise decoupled PPO algorithm that reduces variance in importance ratios by introducing per-token proximal policy logprobs (π_proximal_t) to replace global π_proximal in the AReaL objective function. This enables fine-grained importance sampling correction when different sequence segments are generated under different policy versions during asynchronous training. **Core Algorithm**: - Original: importance_ratio = π_proximal / π_behave - Segment-wise: importance_ratio[i] = π_proximal_t[i] / π_behave[i] - Reduces variance by tracking per-token policy versions during generation - Supports abort-resume with proximal_t updates when policy changes mid-generation **Key Components**: 1. **Generation Engine** (areal/engine/sglang_remote.py): - Track per-token policy versions during generation - Maintain proximal_logprobs_t aligned with output tokens - Recompute proximal_t for previous tokens after abort-resume - Add recompute_output_logprobs_sync() API for cache recomputation 2. **PPO Loss** (realhf/impl/model/utils/ppo_functional.py): - Add behav_imp_weight_floor parameter for symmetric/asymmetric clipping - Support both cap and floor for behav_imp_weight filtering - Compute segment-wise importance weights: exp(proximal_t - old_logprobs) 3. **Actor Interface** (areal/engine/ppo/actor.py): - Roll old_logp to align with rolled proximal_t and prox_logp - Fix tensor misalignment bug in loss computation 4. **Workflow API** (areal/api/workflow_api.py): - Add recompute_all_proximal_t() to update cached samples before weight sync - Track recompute version (RECOMPUTE_VERSION_KEY) to avoid redundant work - Thread-safe queue recompute with drain-putback pattern 5. **Model API** (realhf/api/core/model_api.py): - Add proximal_logprobs_t field to ModelResponse - Extend version tracking for per-token policy versions **Critical Bug Fixes**: Bug #1: Incorrect proximal_logprobs_t length (CRITICAL) - **Issue**: Extended with prompt logprobs and duplicated during abort-resume - **Impact**: Length mismatch causing crashes (proximal_t.len != output_tokens.len) - **Fix**: Changed EXTEND to REPLACE logic for abort-resume iterations - **Location**: areal/engine/sglang_remote.py:240-282 Bug #2: Misaligned tensor roll (CRITICAL) - **Issue**: proximal_t rolled but old_logp not rolled before loss computation - **Impact**: Off-by-one misalignment causing wrong importance weights - **Fix**: Added roll operation for old_logp before ppo_actor_loss_fn() - **Location**: areal/engine/ppo/actor.py:322-327 **Test Coverage** (19 CPU-only tests, all passing): Core functionality: - test_behav_imp_weight_floor.py (10 tests): Cap/floor clipping, symmetric/asymmetric - test_grpo_loss_logp_alignment.py (4 tests): Tensor rolling behavior, alignment - test_proximal_t_generation_simple.py (5 tests): Generation, abort-resume, integration Workflow tests (require einops): - test_recompute_timing.py (11 tests): Cache/queue recompute, drain-putback - test_workflow_executor_wait.py (26 tests): Wait behavior, version purge, caching - test_workflow_helper_functions.py: Helper function unit tests - test_workflow_wait_bugfixes.py: Edge cases and concurrency **Files Modified** (29 files): - Core implementation: 8 files (engine, loss, workflow, API) - Tests: 7 new test files - Documentation: CHANGES_SUMMARY.md, PROXIMAL_T_BUGS_COMPLETE_ANALYSIS.md - Evaluation: eval scripts and configs - Examples: gsm8k_grpo.py, gsm8k_grpo.yaml **Design Decisions**: - Per-token version tracking for fine-grained importance sampling - Abort-resume pattern for policy updates during generation - Recompute-before-sync to ensure proximal_t freshness - Symmetric/asymmetric clipping via cap+floor parameters - Thread-safe drain-putback for queue recompute **Validation**: - All 19 CPU tests pass ✓ - Integration tested with gsm8k_grpo.py (requires GPU) - Documented in docs/PROXIMAL_T_BUGS_COMPLETE_ANALYSIS.md **References**: - AReaL paper (arXiv:2505.24298): Decoupled PPO objective - This implementation: Extends to segment-wise per-token tracking Co-Authored-By: Claude <noreply@anthropic.com>
zhshgmail
pushed a commit
that referenced
this pull request
Oct 10, 2025
Implement segment-wise decoupled PPO algorithm that reduces variance in importance ratios by introducing per-token proximal policy logprobs (π_proximal_t) to replace global π_proximal in the AReaL objective function. This enables fine-grained importance sampling correction when different sequence segments are generated under different policy versions during asynchronous training. **Core Algorithm**: - Original: importance_ratio = π_proximal / π_behave - Segment-wise: importance_ratio[i] = π_proximal_t[i] / π_behave[i] - Reduces variance by tracking per-token policy versions during generation - Supports abort-resume with proximal_t updates when policy changes mid-generation **Key Components**: 1. **Generation Engine** (areal/engine/sglang_remote.py): - Track per-token policy versions during generation - Maintain proximal_logprobs_t aligned with output tokens - Recompute proximal_t for previous tokens after abort-resume - Add recompute_output_logprobs_sync() API for cache recomputation 2. **PPO Loss** (realhf/impl/model/utils/ppo_functional.py): - Add behav_imp_weight_floor parameter for symmetric/asymmetric clipping - Support both cap and floor for behav_imp_weight filtering - Compute segment-wise importance weights: exp(proximal_t - old_logprobs) 3. **Actor Interface** (areal/engine/ppo/actor.py): - Roll old_logp to align with rolled proximal_t and prox_logp - Fix tensor misalignment bug in loss computation 4. **Workflow API** (areal/api/workflow_api.py): - Add recompute_all_proximal_t() to update cached samples before weight sync - Track recompute version (RECOMPUTE_VERSION_KEY) to avoid redundant work - Thread-safe queue recompute with drain-putback pattern 5. **Model API** (realhf/api/core/model_api.py): - Add proximal_logprobs_t field to ModelResponse - Extend version tracking for per-token policy versions **Critical Bug Fixes**: Bug #1: Incorrect proximal_logprobs_t length (CRITICAL) - **Issue**: Extended with prompt logprobs and duplicated during abort-resume - **Impact**: Length mismatch causing crashes (proximal_t.len != output_tokens.len) - **Fix**: Changed EXTEND to REPLACE logic for abort-resume iterations - **Location**: areal/engine/sglang_remote.py:240-282 Bug #2: Misaligned tensor roll (CRITICAL) - **Issue**: proximal_t rolled but old_logp not rolled before loss computation - **Impact**: Off-by-one misalignment causing wrong importance weights - **Fix**: Added roll operation for old_logp before ppo_actor_loss_fn() - **Location**: areal/engine/ppo/actor.py:322-327 **Test Coverage** (19 CPU-only tests, all passing): Core functionality: - test_behav_imp_weight_floor.py (10 tests): Cap/floor clipping, symmetric/asymmetric - test_grpo_loss_logp_alignment.py (4 tests): Tensor rolling behavior, alignment - test_proximal_t_generation_simple.py (5 tests): Generation, abort-resume, integration Workflow tests (require einops): - test_recompute_timing.py (11 tests): Cache/queue recompute, drain-putback - test_workflow_executor_wait.py (26 tests): Wait behavior, version purge, caching - test_workflow_helper_functions.py: Helper function unit tests - test_workflow_wait_bugfixes.py: Edge cases and concurrency **Files Modified** (29 files): - Core implementation: 8 files (engine, loss, workflow, API) - Tests: 7 new test files - Documentation: CHANGES_SUMMARY.md, PROXIMAL_T_BUGS_COMPLETE_ANALYSIS.md - Evaluation: eval scripts and configs - Examples: gsm8k_grpo.py, gsm8k_grpo.yaml **Design Decisions**: - Per-token version tracking for fine-grained importance sampling - Abort-resume pattern for policy updates during generation - Recompute-before-sync to ensure proximal_t freshness - Symmetric/asymmetric clipping via cap+floor parameters - Thread-safe drain-putback for queue recompute **Validation**: - All 19 CPU tests pass ✓ - Integration tested with gsm8k_grpo.py (requires GPU) - Documented in docs/PROXIMAL_T_BUGS_COMPLETE_ANALYSIS.md **References**: - AReaL paper (arXiv:2505.24298): Decoupled PPO objective - This implementation: Extends to segment-wise per-token tracking Co-Authored-By: Claude <noreply@anthropic.com>
zhshgmail
pushed a commit
that referenced
this pull request
Oct 10, 2025
Implement segment-wise decoupled PPO algorithm that reduces variance in importance ratios by introducing per-token proximal policy logprobs (π_proximal_t) to replace global π_proximal in the AReaL objective function. This enables fine-grained importance sampling correction when different sequence segments are generated under different policy versions during asynchronous training. **Core Algorithm**: - Original: importance_ratio = π_proximal / π_behave - Segment-wise: importance_ratio[i] = π_proximal_t[i] / π_behave[i] - Reduces variance by tracking per-token policy versions during generation - Supports abort-resume with proximal_t updates when policy changes mid-generation **Key Components**: 1. **Generation Engine** (areal/engine/sglang_remote.py): - Track per-token policy versions during generation - Maintain proximal_logprobs_t aligned with output tokens - Recompute proximal_t for previous tokens after abort-resume - Add recompute_output_logprobs_sync() API for cache recomputation 2. **PPO Loss** (realhf/impl/model/utils/ppo_functional.py): - Add behav_imp_weight_floor parameter for symmetric/asymmetric clipping - Support both cap and floor for behav_imp_weight filtering - Compute segment-wise importance weights: exp(proximal_t - old_logprobs) 3. **Actor Interface** (areal/engine/ppo/actor.py): - Roll old_logp to align with rolled proximal_t and prox_logp - Fix tensor misalignment bug in loss computation 4. **Workflow API** (areal/api/workflow_api.py): - Add recompute_all_proximal_t() to update cached samples before weight sync - Track recompute version (RECOMPUTE_VERSION_KEY) to avoid redundant work - Thread-safe queue recompute with drain-putback pattern 5. **Model API** (realhf/api/core/model_api.py): - Add proximal_logprobs_t field to ModelResponse - Extend version tracking for per-token policy versions **Critical Bug Fixes**: Bug #1: Incorrect proximal_logprobs_t length (CRITICAL) - **Issue**: Extended with prompt logprobs and duplicated during abort-resume - **Impact**: Length mismatch causing crashes (proximal_t.len != output_tokens.len) - **Fix**: Changed EXTEND to REPLACE logic for abort-resume iterations - **Location**: areal/engine/sglang_remote.py:240-282 Bug #2: Misaligned tensor roll (CRITICAL) - **Issue**: proximal_t rolled but old_logp not rolled before loss computation - **Impact**: Off-by-one misalignment causing wrong importance weights - **Fix**: Added roll operation for old_logp before ppo_actor_loss_fn() - **Location**: areal/engine/ppo/actor.py:322-327 **Test Coverage** (19 CPU-only tests, all passing): Core functionality: - test_behav_imp_weight_floor.py (10 tests): Cap/floor clipping, symmetric/asymmetric - test_grpo_loss_logp_alignment.py (4 tests): Tensor rolling behavior, alignment - test_proximal_t_generation_simple.py (5 tests): Generation, abort-resume, integration Workflow tests (require einops): - test_recompute_timing.py (11 tests): Cache/queue recompute, drain-putback - test_workflow_executor_wait.py (26 tests): Wait behavior, version purge, caching - test_workflow_helper_functions.py: Helper function unit tests - test_workflow_wait_bugfixes.py: Edge cases and concurrency **Files Modified** (29 files): - Core implementation: 8 files (engine, loss, workflow, API) - Tests: 7 new test files - Documentation: CHANGES_SUMMARY.md, PROXIMAL_T_BUGS_COMPLETE_ANALYSIS.md - Evaluation: eval scripts and configs - Examples: gsm8k_grpo.py, gsm8k_grpo.yaml **Design Decisions**: - Per-token version tracking for fine-grained importance sampling - Abort-resume pattern for policy updates during generation - Recompute-before-sync to ensure proximal_t freshness - Symmetric/asymmetric clipping via cap+floor parameters - Thread-safe drain-putback for queue recompute **Validation**: - All 19 CPU tests pass ✓ - Integration tested with gsm8k_grpo.py (requires GPU) - Documented in docs/PROXIMAL_T_BUGS_COMPLETE_ANALYSIS.md **References**: - AReaL paper (arXiv:2505.24298): Decoupled PPO objective - This implementation: Extends to segment-wise per-token tracking Co-Authored-By: Claude <noreply@anthropic.com>
zhshgmail
pushed a commit
that referenced
this pull request
Oct 10, 2025
Implement segment-wise decoupled PPO algorithm that reduces variance in importance ratios by introducing per-token proximal policy logprobs (π_proximal_t) to replace global π_proximal in the AReaL objective function. This enables fine-grained importance sampling correction when different sequence segments are generated under different policy versions during asynchronous training. **Core Algorithm**: - Original: importance_ratio = π_proximal / π_behave - Segment-wise: importance_ratio[i] = π_proximal_t[i] / π_behave[i] - Reduces variance by tracking per-token policy versions during generation - Supports abort-resume with proximal_t updates when policy changes mid-generation **Key Components**: 1. **Generation Engine** (areal/engine/sglang_remote.py): - Track per-token policy versions during generation - Maintain proximal_logprobs_t aligned with output tokens - Recompute proximal_t for previous tokens after abort-resume - Add recompute_output_logprobs_sync() API for cache recomputation 2. **PPO Loss** (realhf/impl/model/utils/ppo_functional.py): - Add behav_imp_weight_floor parameter for symmetric/asymmetric clipping - Support both cap and floor for behav_imp_weight filtering - Compute segment-wise importance weights: exp(proximal_t - old_logprobs) 3. **Actor Interface** (areal/engine/ppo/actor.py): - Roll old_logp to align with rolled proximal_t and prox_logp - Fix tensor misalignment bug in loss computation 4. **Workflow API** (areal/api/workflow_api.py): - Add recompute_all_proximal_t() to update cached samples before weight sync - Track recompute version (RECOMPUTE_VERSION_KEY) to avoid redundant work - Thread-safe queue recompute with drain-putback pattern 5. **Model API** (realhf/api/core/model_api.py): - Add proximal_logprobs_t field to ModelResponse - Extend version tracking for per-token policy versions **Critical Bug Fixes**: Bug #1: Incorrect proximal_logprobs_t length (CRITICAL) - **Issue**: Extended with prompt logprobs and duplicated during abort-resume - **Impact**: Length mismatch causing crashes (proximal_t.len != output_tokens.len) - **Fix**: Changed EXTEND to REPLACE logic for abort-resume iterations - **Location**: areal/engine/sglang_remote.py:240-282 Bug #2: Misaligned tensor roll (CRITICAL) - **Issue**: proximal_t rolled but old_logp not rolled before loss computation - **Impact**: Off-by-one misalignment causing wrong importance weights - **Fix**: Added roll operation for old_logp before ppo_actor_loss_fn() - **Location**: areal/engine/ppo/actor.py:322-327 **Test Coverage** (19 CPU-only tests, all passing): Core functionality: - test_behav_imp_weight_floor.py (10 tests): Cap/floor clipping, symmetric/asymmetric - test_grpo_loss_logp_alignment.py (4 tests): Tensor rolling behavior, alignment - test_proximal_t_generation_simple.py (5 tests): Generation, abort-resume, integration Workflow tests (require einops): - test_recompute_timing.py (11 tests): Cache/queue recompute, drain-putback - test_workflow_executor_wait.py (26 tests): Wait behavior, version purge, caching - test_workflow_helper_functions.py: Helper function unit tests - test_workflow_wait_bugfixes.py: Edge cases and concurrency **Files Modified** (29 files): - Core implementation: 8 files (engine, loss, workflow, API) - Tests: 7 new test files - Documentation: CHANGES_SUMMARY.md, PROXIMAL_T_BUGS_COMPLETE_ANALYSIS.md - Evaluation: eval scripts and configs - Examples: gsm8k_grpo.py, gsm8k_grpo.yaml **Design Decisions**: - Per-token version tracking for fine-grained importance sampling - Abort-resume pattern for policy updates during generation - Recompute-before-sync to ensure proximal_t freshness - Symmetric/asymmetric clipping via cap+floor parameters - Thread-safe drain-putback for queue recompute **Validation**: - All 19 CPU tests pass ✓ - Integration tested with gsm8k_grpo.py (requires GPU) - Documented in docs/PROXIMAL_T_BUGS_COMPLETE_ANALYSIS.md **References**: - AReaL paper (arXiv:2505.24298): Decoupled PPO objective - This implementation: Extends to segment-wise per-token tracking Co-Authored-By: Claude <noreply@anthropic.com>
zhshgmail
added a commit
that referenced
this pull request
Nov 4, 2025
…eal-project#501) * feat: implement GSPO (Group-level Sequential Policy Optimization) Add GSPO support with sequence-level advantage computation instead of token-level. This implementation includes: - Add importance_sampling_level configuration ("token" for PPO, "sequence" for GSPO) - Modify ppo_actor_loss_fn to compute advantages at sequence level for GSPO - Add comprehensive tests for sequence-level importance sampling - Add GSPO example with GSM8K dataset configuration - Update weight_update_meta initialization and validation workflow - Merge DAPO configs for better flexibility - Update README with GSPO usage and configuration details The key difference from standard PPO is that GSPO sums token-level advantages within each sequence before computing the policy gradient, enabling more stable sequence-level optimization. * Address PR review comments Changes: - Add more assertions in unit tests for GSPO - Fix GSPO packed sequences to use cu_seqlens for per-sequence ratios - Fix GSPO gradient scaling bug by averaging instead of summing advantages - Add GSPO algorithm description and implementation notes to documentation - Add PPOActor configuration logging for reproducibility - Simplify documentation by removing detailed implementation notes - Update examples to reuse existing gsm8k_grpo.py with GSPO config - Remove dedicated gsm8k_gspo.py script (not necessary) - Remove unnecessary overlong penalty parameters from config GSPO can be enabled with any PPO-like training by setting importance_sampling_level: sequence in the configuration. * Ajust code format to pass pipeline black formatter check. * perf: vectorize GSPO 1D packed tensor computation and simplify docs - Replace Python for loop with scatter_add_ for GPU performance - Simplify advantage aggregation description in documentation - All 18 GSPO tests pass with identical results Addresses GitHub review comments #1 and #2. * Ajust code format to pass pipeline black formatter check. --------- Co-authored-by: bruceli_ed24 <bruce.li2@h-partners.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
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.
This PR adds a comprehensive PlantUML diagram that illustrates the high-level architecture of the AReaL (Ant Reasoning Reinforcement Learning) system, addressing the need for clear documentation of the system's innovative asynchronous RL training architecture.
What's Added
Core Documentation
docs/architecture/system-overview.puml- Complete PlantUML source diagram showing the system architecturedocs/architecture/system-overview.png- Generated visual diagram (256KB)docs/architecture/README.md- Comprehensive documentation explaining the diagram and its usageArchitecture Coverage
The diagram comprehensively covers all major system components:
API Layer (
realhf/api/): CLI arguments, core APIs (System, Model, Data, Agent, Environment), and quickstart functionalityModel Implementation (
realhf/impl/model/): Backend integrations including SGLang, vLLM, and Megatron, neural network modules (attention, MLP, embeddings, MOE), and parallelism support (pipeline parallel, tensor parallel, parameter reallocation)Training System (
realhf/system/): Core workers (Model Worker, Rollout Worker, Generation Server), distributed control systems, data management components, and communication infrastructureInfrastructure (
realhf/base/): Logging, name resolution, topology management, GPU utilities, monitoring, and networkingScheduler (
realhf/scheduler/): Job scheduling clients, evaluators, and execution backends (local/SLURM)Implementation & Experiments: Concrete implementations for agents, datasets, environments, and both asynchronous and traditional experiment configurations
Key Architectural Insights
The diagram specifically highlights AReaL's fully asynchronous RL training architecture, showing:
Benefits
The diagram is generated using industry-standard PlantUML and can be easily updated as the system evolves. It successfully captures the essence of AReaL's breakthrough asynchronous reinforcement learning architecture for large language models.
Warning
Firewall rules blocked me from connecting to one or more addresses
I tried to connect to the following addresses, but was blocked by firewall rules:
esm.ubuntu.com/usr/lib/apt/methods/https(dns block)If you need me to access, download, or install something from one of these locations, you can either:
This pull request was created as a result of the following prompt from Copilot chat.
💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.