fix: initialize stage_indices in deepseek policy for interleave pipeline

[AmSach]

Fix: Missing stage_indices in DeepSeek Policy (Issue hpcaitech#6397)

Repo: hpcaitech/ColossalAI
Issue: hpcaitech#6397 — AttributeError: PipelineStageManager object has no attribute stage_indices

What was broken

When running DeepSeek with interleaved pipeline parallelism (--pp_style interleaved --n_chunks 2), the code crashes at interleaved_pp.py:309:

internal_inputs["stage_index"] = self.stage_manager.stage_indices[model_chunk_id]

The root cause: in colossalai/shardformer/policies/deepseek.py, the get_held_layers() method's interleave branch never initializes stage_manager.stage_indices. The same code path in deepseek_v3.py (line 120) does initialize it correctly.

How I fixed it

Added one line to get_held_layers() in deepseek.py:

stage_indices = stage_manager.get_stage_index(layers_per_stage)
stage_manager.stage_indices = stage_indices  # ← added this line

Diff

stage_indices = stage_manager.get_stage_index(layers_per_stage)

• stage_manager.stage_indices = stage_indices
  if stage_manager.is_first_stage(ignore_chunk=True):

Verification

• Syntax check: passed (verified with ast.parse)
• Mirrors the working deepseek_v3.py implementation at line 120

Closes hpcaitech#6397