Skip to content

Add comprehensive PlantUML system architecture diagram for AReaL#2

Closed
zhshgmail with Copilot wants to merge 1 commit into
mainfrom
copilot/fix-654db62d-78c1-480e-8e16-0937d16be242
Closed

Add comprehensive PlantUML system architecture diagram for AReaL#2
zhshgmail with Copilot wants to merge 1 commit into
mainfrom
copilot/fix-654db62d-78c1-480e-8e16-0937d16be242

Conversation

Copilot AI commented Jul 29, 2025

Copy link
Copy Markdown

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 architecture
  • docs/architecture/system-overview.png - Generated visual diagram (256KB)
  • docs/architecture/README.md - Comprehensive documentation explaining the diagram and its usage

Architecture Coverage

The diagram comprehensively covers all major system components:

API Layer (realhf/api/): CLI arguments, core APIs (System, Model, Data, Agent, Environment), and quickstart functionality

Model 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 infrastructure

Infrastructure (realhf/base/): Logging, name resolution, topology management, GPU utilities, monitoring, and networking

Scheduler (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:

  1. Decoupled Generation and Training: How the Generation Server operates continuously without waiting for training completion
  2. Asynchronous Data Flow: The streaming architecture that enables maximum throughput through Push/Pull streams and buffering
  3. Distributed Control: How the Controller manages generation and training workers independently
  4. Component Relationships: Clear visualization of how different system layers interact

Benefits

  • Developer Onboarding: New contributors can quickly understand the system architecture
  • System Understanding: Visualizes the innovative async RL design that achieves 2.77x speedup
  • Documentation: Provides a living document that reflects the actual codebase structure
  • Reference: Serves as a comprehensive overview for both technical and non-technical stakeholders

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
    • Triggering command: /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.

Create a comprehensive PlantUML diagram that illustrates the high-level architecture of the AReaL (Ant Reasoning Reinforcement Learning) system. The diagram should include:

Requirements:

  1. Core Components: Show the main system components including:

    • API layer (CLI args, core APIs, quickstart)
    • Model implementations (backends like SGLang, modules, parallelism)
    • Training systems (schedulers, experiments)
    • Base infrastructure (logging, topology, name resolution)
    • Evaluation systems
  2. System Architecture: Visualize the decoupled generation and training architecture that AReaL implements for RL training of LLMs

  3. Key Relationships: Show how different components interact, particularly:

    • The relationship between model backends and training systems
    • How the scheduler manages distributed training
    • The flow from API to implementation layers
  4. File Organization: The diagram should reflect the actual code structure found in the repository, including the main realhf/ package structure

Implementation Details:

  • Use PlantUML component diagram syntax
  • Include appropriate groupings and packages to show system layers
  • Add brief descriptions for major components
  • Use consistent styling and clear component relationships
  • Save as docs/architecture/system-overview.puml

This diagram will help developers and users understand the overall system architecture and how different parts of AReaL work together to provide reinforcement learning capabilities for large language models.


💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.

@zhshgmail zhshgmail marked this pull request as ready for review July 29, 2025 04:42
@zhshgmail zhshgmail closed this Jul 29, 2025
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
Copilot AI requested a review from zhshgmail July 29, 2025 05:11
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
- 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.
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>
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.

2 participants