What causes "Pi05 action state is missing on single-rank run" in sgl-project/sglang?

Run through the normal sglang serving/scheduler path so the distributed group and action-state plumbing is initialized

sgl-project/sglang · error · RuntimeError

Pi05 action state is missing on single-rank run

Error message

Pi05 action state is missing on single-rank run

What it means

During sample_actions with action sequence-parallelism configured, the initial noisy action tensor x_t is normally created on the action-root rank and broadcast to the group. On a single-rank run (get_vla_split_group() returns None) there is no broadcast, so the caller must supply x_t; if it is None the policy cannot proceed and raises. It indicates a code-path inconsistency rather than a user config problem.

Source

Thrown at python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py:1082

            "runtime_role": self.runtime_role,
            "world_size": split.group.world_size,
            "prefix_root": split.prefix_root,
            "action_root": split.action_root,
            "action_ranks": list(split.action_ranks),
            "action_sequence_parallel": self._can_use_action_sequence_parallel(
                prefix_context,
                self.config.action_horizon,
            ),
        }

    def _broadcast_initial_action_state(
        self,
        x_t: torch.Tensor | None,
    ) -> torch.Tensor:
        split = get_vla_split_group()
        if split is None:
            if x_t is None:
                raise RuntimeError("Pi05 action state is missing on single-rank run")
            return x_t
        x_t = broadcast_tensor_from_rank(
            x_t,
            split,
            src=split.action_root,
            device=self.device,
        )
        if x_t is None:
            raise RuntimeError("Pi05 action state broadcast returned None")
        return x_t

    def _shard_action_sequence(self, x_t: torch.Tensor) -> tuple[torch.Tensor, int]:
        sp_world_size = get_sequence_parallel_world_size()
        sp_rank = get_sp_parallel_rank()
        local_len = x_t.shape[1] // sp_world_size
        start = sp_rank * local_len
        end = start + local_len
        return x_t[:, start:end].contiguous(), start

View on GitHub (pinned to 0132848349)

Solutions

Run through the normal sglang serving/scheduler path so the distributed group and action-state plumbing is initialized
If calling sample_actions directly, pass an explicit x_t tensor of the right shape (float32, [batch, action_len, action_dim])
Update sglang — this path mismatch may be a fixed bug
Check that sequence-parallel flags are consistent: either fully enable action SP or fully disable it

Example fix

# before
actions = policy.sample_actions(prefix, suffix, x_t=None)  # single-rank, raises

# after
x_t = torch.randn(batch, action_len, action_dim, device=policy.device, dtype=torch.float32)
actions = policy.sample_actions(prefix, suffix, x_t=x_t)

Defensive patterns

Strategy: fallback

Validate before calling

from sglang.multimodal_gen.runtime.models.vlas.pi05_policy import get_vla_split_group
if get_vla_split_group() is None and x_t is None and not i_own_noise_init:
    x_t = torch.randn(batch, action_len, action_dim, device=device, dtype=torch.float32)

Try / catch

try:
    actions = policy.sample_actions(prefix, suffix, x_t=x_t)
except RuntimeError as e:
    if "action state is missing on single-rank run" in str(e):
        x_t = torch.randn(batch, action_len, action_dim,
                          device=policy.device, dtype=torch.float32)
        actions = policy.sample_actions(prefix, suffix, x_t=x_t)
    else:
        raise

Prevention

When calling sample_actions outside the scheduler, always pass an explicit x_t
Initialize the distributed/VLA split runtime through the standard sglang launcher rather than ad-hoc scripts

When it happens

Trigger: Calling sample_actions with x_t=None (the normal fallback path where the policy generates its own noise) while the run is single-rank / the VLA split group is not initialized — i.e. the internal condition that should have populated x_t did not fire.

Common situations: Running Pi05 inference without the distributed/sequence-parallel runtime initialized, or an sglang version mismatch where the fallback noise-initialization branch was removed or guarded differently; calling sample_actions directly in unit tests outside the scheduler.

Related errors

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/3d6501f7da84d145. Report an issue: GitHub.