Pull Request: Auto-Parallelization and Topology-Aware Node Selection for mindspeed_rl

Summary

This PR adds automatic parallelization configuration for RLHF workloads and rack-aware topology selection for better resource allocation in multi-rack environments. These features enable the system to automatically determine optimal Tensor Parallelism (TP), Data Parallelism (DP), Pipeline Parallelism (PP), Context Parallelism (CP), Expert Parallelism (EP), and Sequence Parallelism (SP) configurations for both inference and training, and intelligently place workers based on cluster topology.

Features Added

1. Auto-Parallelization for Inference and Training

Automatically determines optimal TP/CP/DP/EP/SP/PP configuration for inference and training based on:

• Available hardware resources (number of NPUs)
• Model constraints (hidden dimension, attention heads, sequence length, number of layers, number of experts)
• Runtime memory checking (loads model to verify memory fit)

Key Benefits:

• Memory-aware: Uses runtime memory checks to avoid OOM errors (more reliable than analytical formulas)
• Follows best practices: Aligns with Megatron-LM recommendations for parallelism strategies
• MOE support: Handles Expert Parallelism correctly (EP happens inside DP, not a separate dimension)

2. Rack-Aware Topology Selection

Intelligent node allocation that prioritizes:

• Single-rack allocation when possible (better performance)
• Symmetric allocation across racks when single-rack is not possible
• TP/PP/CP/EP group constraints respected within rack boundaries
• Balanced resource distribution across racks

Balanced Node Selection:

• select_balanced_nodes() implements intelligent placement rules:
  1. Single-rack preference: Allocates all nodes from one rack if capacity allows
  2. Symmetric allocation: Distributes nodes evenly across multiple racks when needed
  3. Group constraints: Respects TP/PP/CP/EP group sizes (keeps groups within same rack when possible)
  4. TP-aware: Ensures TP groups are placed optimally for low-latency communication

Files Changed

Core Implementation

• mindspeed_rl/workers/scheduler/launcher.py:
  • auto_parallelize_inference(): Main entry point for inference auto-parallelization
  • auto_parallelize_inference_simple(): Core algorithm for TP/CP/EP selection (inference)
  • auto_parallelize_training(): Core algorithm for TP/PP/CP/EP/SP selection (training)
  • check_inference_parallel_constraints(): Validates TP/CP against model constraints
  • check_training_parallel_constraints(): Validates TP/PP/CP against model constraints
  • check_memory_at_runtime(): Runtime memory checking via model initialization
  • estimate_memory_requirements(): Analytical memory estimation (fallback, with warnings)
  • find_factors(): Helper function to find all factors of a number (for PP/CP candidates)
  • get_cluster_topology(): Detects rack topology from Ray node resources (RACK_X)
  • select_balanced_nodes(): Implements balanced node selection algorithm with single-rack preference and symmetric allocation
  • get_rack_capacities(): Calculates capacity per rack for allocation decisions
  • _select_nodes_group_within_rack(): Allocates nodes keeping TP/PP/CP groups within same rack when possible

Configuration

• mindspeed_rl/config_cls/generate_config.py:

  • Added auto_parallel (bool): Enable/disable auto-parallelization for inference
  • Added auto_parallel_check_memory (bool): Use runtime checks vs analytical estimation

• mindspeed_rl/config_cls/megatron_config.py:

  • Added auto_parallel (bool): Enable/disable auto-parallelization for training

• mindspeed_rl/config_cls/rl_config.py:

  • Added auto_parallel attribute for backwards compatibility (ignored)

Integration

• cli/train_grpo.py:
  • Integrated auto-parallelization into training workflow
  • Reads auto_parallel flag from actor_config for training
  • Calls auto_parallelize_training() when enabled for actor/ref/reward configs
  • Calls auto_parallelize_inference() when enabled in generate_config

Documentation

• docs/AUTO_PARALLEL_IMPLEMENTATION.md: Comprehensive documentation
• docs/DEVICE_MEMORY_DETECTION.md: Device memory detection mechanism

Configuration

YAML Configuration

# Inference auto-parallelization
generate_config:
  # Enable auto-parallelization for inference
  auto_parallel: true
  
  # Runtime memory checking (default: true, recommended)
  auto_parallel_check_memory: true  # true = reliable, false = fast but unreliable
  
  # These will be auto-selected when auto_parallel=true
  infer_tensor_parallel_size: 2  # Will be overridden
  infer_pipeline_parallel_size: 1 # Will not be overridden (vLLM doesn't support PP)
  infer_expert_parallel_size: 1   # Will be auto-selected for MOE models

# Training auto-parallelization
actor_config:  # Same for ref_config and reward_config
  # Enable auto-parallelization for training
  auto_parallel: true
  
  # These will be auto-selected when auto_parallel=true
  tensor_model_parallel_size: 1      # Will be overridden
  pipeline_model_parallel_size: 1    # Will be auto-selected
  context_parallel_size: 1           # Will be auto-selected
  expert_model_parallel_size: 1      # Will be auto-selected for MOE models
  sequence_parallel: false           # Will be auto-enabled when needed (see below)

Usage Examples

Basic Usage

# config.yaml
generate_config:
  auto_parallel: true
  auto_parallel_check_memory: true

The system will automatically:

• For inference: Find minimum TP that satisfies constraints and fits in memory, select optimal CP if sequence length > 16000, determine EP for MOE models
• For training: Find optimal TP/PP/CP combination that minimizes (TPPPCP) to maximize DP, determine EP for MOE models, automatically enable SP when needed
• Update config with selected values (TP, PP, CP, EP, SP)

With Rack Configuration

# Set rack information via Ray resources on each node
ray start --resources='{"RACK_0": 1.0}'  # Node in rack 0
# On another node:
ray start --resources='{"RACK_1": 1.0}'  # Node in rack 1

The system will:

1. Automatically detect rack topology via get_cluster_topology()
2. Use select_balanced_nodes() for intelligent placement:
   • Prefer single-rack allocation when possible
   • Use symmetric allocation across racks when needed
   • Respect TP/PP/CP/EP group constraints
3. Create placement groups with rack-aware constraints

If no RACK_X resources are found, falls back to sequential placement without rack awareness.

Testing

Test Configuration

Tested with:

• Model: Qwen2.5-3B
• Hardware: Ascend NPU (A2)
• NPUs: 16 per test case

Test Result:

• Topology search improvement:

1. Two servers on different racks: (E2E time 289 sec)

2. Two servers on Same racks:(E2E time 276 sec)

3. auto-parallel improvement:(E2E time 243 sec) (No topology search)

Improvement for Topology search: 5%

Improvement for Auto-parallel: 18%

Breaking Changes

None - All changes are backward compatible:

• Auto-parallelization is disabled by default (auto_parallel: false)
• Existing configs continue to work without changes
• New config fields have safe defaults

Performance Considerations

Runtime Memory Checks (Default)

• Accuracy: High - actually loads model to verify memory
• Speed: Slower (loads model per TP candidate, typically seconds)
• Recommendation: Use for production deployments

Analytical Estimation (Optional)

• Accuracy: Low - formulas miss many production factors
• Speed: Fast (instant, no model loading)
• Recommendation: Only for development/testing with warnings

Memory Detection

The system automatically detects device memory:

1. CUDA: Uses torch.cuda.get_device_properties().total_memory
2. NPU: Tries torch.npu.get_device_properties() or torch.npu.memory_info()
3. Environment: Falls back to NPU_MEMORY_GB environment variable
4. Default: Conservative 16GB default if detection fails

See docs/DEVICE_MEMORY_DETECTION.md for details.

Documentation

Comprehensive documentation provided:

• docs/AUTO_PARALLEL_IMPLEMENTATION.md: Complete implementation guide
• docs/DEVICE_MEMORY_DETECTION.md: Memory detection mechanisms
• Inline code comments and docstrings

Technical Details

Future Enhancements

Potential improvements (not in this PR):

• Performance profiling-based selection
• Distributed memory checking across nodes

Request for Review and Merge

We request the MindSpeed-RL team to review and merge this PR. The changes:

Review Focus Areas

Please pay special attention to:

1. Memory checking accuracy: Runtime checks vs analytical estimation trade-offs
2. Topology detection: Detection of RACK_X resources from Ray
3. Balanced node selection: select_balanced_nodes() algorithm (single-rack preference, symmetric allocation)
4. Constraint validation: TP/PP/CP/EP constraint checking algorithms
5. EP treatment: EP as part of DP vs separate dimension (we use EP inside DP)
6. SP auto-enablement: Logic for enabling SP based on TP and EP values
7. PP selection strategy: Minimize (TPPPCP) to maximize DP, increase PP first when TP=1 doesn't fit. PP is not supported by vllm now.
8. Integration points: Integration with existing training workflows

Questions or Concerns

If you have questions or concerns, please:

• Comment on specific lines in the code
• Request clarification on algorithm choices
• Suggest improvements or alternatives

We're open to feedback and ready to make adjustments based on your review.

---

**Thank you for your time and con