minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

A runner and the objects it builds freeze the placement they describe (#40341)

Browse Source
This commit is contained in:
Cheng Wan
2026-09-21 12:24:17 -07:00
committed by GitHub
parent 2d0e94e3a3
commit 65be3fa71a
13 changed files with 289 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/bootstrap.py
+1 -19
View File
@@ -3,14 +3,12 @@ import os
import time
from typing import List, Optional
import msgspec
import torch
import torch.distributed as dist
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed import (
get_default_distributed_backend,
get_pp_group,
get_tp_group,
get_world_group,
init_distributed_environment,
@@ -60,13 +58,6 @@ _is_cpu_arm64 = is_host_cpu_arm64()
_TP_ALL_TO_ALL_WARMUP_BYTES_PER_PEER = 4 << 20
class TorchDistributedResult(msgspec.Struct, frozen=True, kw_only=True):
tp_group: object
pp_group: object
attention_tp_group: object
pre_model_load_memory: float
def init_torch_distributed(
*,
server_args: ServerArgs,
@@ -140,10 +131,6 @@ def init_torch_distributed(
distributed=get_world_group().world_size > 1 and not is_draft_worker,
cpu_group=get_world_group().cpu_group,
)
tp_group = get_tp_group()
pp_group = get_pp_group()
attention_tp_group = get_parallel().attn_tp_group
# Check memory for tensor parallelism
local_gpu_memory = get_available_gpu_memory(device, ps.gpu_id)
if ps.tp_size > 1 and not is_draft_worker:
@@ -156,12 +143,7 @@ def init_torch_distributed(
f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, "
f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB"
)
return TorchDistributedResult(
tp_group=tp_group,
pp_group=pp_group,
attention_tp_group=attention_tp_group,
pre_model_load_memory=pre_model_load_memory,
)
return pre_model_load_memory
def _resolve_backend(*, device: str) -> str:
python/sglang/srt/managers/scheduler.py
-1
View File
@@ -601,7 +601,6 @@ class Scheduler(
and get_parallel().attn_tp_rank == 0
and get_parallel().attn_cp_rank == 0
),
ps=self.ps,
tp_group=self.tp_group,
pp_group=self.pp_group,
enable_hierarchical_cache=self.enable_hierarchical_cache,
python/sglang/srt/mem_cache/kv_cache_builder.py
+9 -8
View File
@@ -64,7 +64,6 @@ if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed.parallel_state import GroupCoordinator
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.managers.tp_worker import BaseTpWorker
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.base_spec_worker import HiCacheDraftPlan
@@ -263,12 +262,14 @@ def build_kv_cache(
attn_cp_cpu_group: ProcessGroup,
enable_metrics: bool,
enable_kv_cache_events: bool,
ps: ParallelState,
tp_group: GroupCoordinator,
pp_group: GroupCoordinator,
enable_hierarchical_cache: bool,
hicache_draft_plan: Optional[HiCacheDraftPlan] = None,
) -> KVCacheBuildResult:
# Built from the scheduler loop, outside any draft scope, so the context
# answers for the process this cache belongs to.
parallel = get_parallel()
sliding_window_size: Optional[int] = None
full_tokens_per_layer: Optional[int] = None
swa_tokens_per_layer: Optional[int] = None
@@ -369,10 +370,10 @@ def build_kv_cache(
enable_session_radix_cache=get_memory().enable_session_radix_cache,
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy,
pp_rank=ps.pp_rank,
pp_size=ps.pp_size,
attn_cp_rank=ps.attn_cp_rank,
attn_cp_size=ps.attn_cp_size,
pp_rank=parallel.pp_rank,
pp_size=parallel.pp_size,
attn_cp_rank=parallel.attn_cp_rank,
attn_cp_size=parallel.attn_cp_size,
chunked_prefill_size=effective_chunked_prefill_size,
sliding_window_size=sliding_window_size,
mtp_draft_device_pools=mtp_draft_device_pools,
@@ -390,8 +391,8 @@ def build_kv_cache(
effective_chunked_prefill_size=effective_chunked_prefill_size,
tp_worker=tp_worker,
model_config=model_config,
tp_size=ps.tp_size,
tp_rank=ps.tp_rank,
tp_size=parallel.tp_size,
tp_rank=parallel.tp_rank,
tp_group=tp_group,
)
with auto_size_hicache(
python/sglang/srt/mem_cache/kv_cache_configurator.py
+13 -11
View File
@@ -206,7 +206,6 @@ def _pp_local_per_request_bytes(
if TYPE_CHECKING:
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.mem_cache.unified_memory_pool import (
UnifiedKVPool,
UnifiedPoolBundle,
@@ -260,7 +259,10 @@ class _PoolSizes(msgspec.Struct, frozen=True, kw_only=True):
class KVCacheConfigurator:
device: str
gpu_id: int
ps: ParallelState
# Frozen at construction, not asked for later: this configurator is built
# inside the scope that describes a draft runner and used outside it.
attn_dp_size: int
pp_size: int
pp_group: Any
model: Any
model_config: ModelConfig
@@ -2324,7 +2326,7 @@ class KVCacheConfigurator:
max_num_reqs = get_schedule().max_running_requests
if max_num_reqs is not None:
requested_per_worker = max_num_reqs // self.ps.attn_dp_size
requested_per_worker = max_num_reqs // self.attn_dp_size
max_num_reqs = min(requested_per_worker, token_capacity // 2)
else:
requested_per_worker = None
@@ -2433,16 +2435,16 @@ class KVCacheConfigurator:
# allocates its own [start_layer, end_layer) slice. Charge the largest
# per-stage share so every rank derives the same pool without a collective.
all_mamba_layers = config.mamba2_cache_params.layers
if self.ps.pp_size > 1 and all_mamba_layers:
if self.pp_size > 1 and all_mamba_layers:
max_stage_mamba_layers = max(
sum(1 for i in all_mamba_layers if start <= i < end)
for start, end in (
get_pp_indices(
self.model_config.num_hidden_layers,
rank,
self.ps.pp_size,
self.pp_size,
)
for rank in range(self.ps.pp_size)
for rank in range(self.pp_size)
)
)
else:
@@ -2474,7 +2476,7 @@ class KVCacheConfigurator:
replayssm_ring_per_req = int(replayssm_ring_per_req * pp_layer_scale)
if replayssm_active and self.hybrid_kda_config is None:
replay_req_slots = (
get_schedule().max_running_requests // self.ps.attn_dp_size + 1
get_schedule().max_running_requests // self.attn_dp_size + 1
)
replayssm_fixed_bytes = replayssm_ring_per_req * replay_req_slots
replayssm_ring_per_slot = 0
@@ -2490,7 +2492,7 @@ class KVCacheConfigurator:
get_context().override(
"mamba_pool.per_dp_shard",
max_mamba_cache_size=get_schedule().max_mamba_cache_size
// self.ps.attn_dp_size,
// self.attn_dp_size,
)
# Reserve intermediate memory based on capped max_num_reqs (+1: the
# pool's padding slot, see memory_pool.py). Skipped under replayssm
@@ -2498,7 +2500,7 @@ class KVCacheConfigurator:
if has_spec_dec and not replayssm_active:
ratio = self._calculate_mamba_ratio()
capped_reqs = min(
get_schedule().max_running_requests // self.ps.attn_dp_size,
get_schedule().max_running_requests // self.attn_dp_size,
get_schedule().max_mamba_cache_size // ratio,
)
intermediate_size = (
@@ -2515,7 +2517,7 @@ class KVCacheConfigurator:
get_context().override(
"mamba_pool.from_max_running_requests",
max_mamba_cache_size=get_schedule().max_running_requests
// self.ps.attn_dp_size,
// self.attn_dp_size,
)
# Reserve intermediate memory based on capped max_num_reqs (+1: the
# pool's padding slot). Skipped under replayssm.
@@ -2555,7 +2557,7 @@ class KVCacheConfigurator:
# Intermediate memory is included in mamba_budget, subtract it
# so the return value only has main_state subtracted from total
capped_reqs = min(
get_schedule().max_running_requests // self.ps.attn_dp_size,
get_schedule().max_running_requests // self.attn_dp_size,
get_schedule().max_mamba_cache_size // ratio,
)
intermediate_size = per_req * (capped_reqs + 1) * D
python/sglang/srt/model_executor/model_runner.py
+59 -42
View File
@@ -421,7 +421,7 @@ class ModelRunner:
import os
logger.warning(
f"Context: {self.device=} {ps.gpu_id=} {os.environ.get('CUDA_VISIBLE_DEVICES')=} {ps.tp_rank=} {ps.tp_size=}"
f"Context: {self.device=} {ps.gpu_id=} {os.environ.get('CUDA_VISIBLE_DEVICES')=} {get_parallel().tp_rank=} {get_parallel().tp_size=}"
)
raise
@@ -447,7 +447,7 @@ class ModelRunner:
# CPU offload
set_offloader(create_offloader(dp_rank=get_parallel().dp_rank))
self._weight_checker = WeightChecker(get_model=lambda: self.model, ps=self.ps)
self._weight_checker = WeightChecker(get_model=lambda: self.model)
if envs.SGLANG_DETECT_SLOW_RANK.get():
slow_rank_detector.execute()
@@ -503,9 +503,9 @@ class ModelRunner:
):
return
join_effective_ep_size = get_parallel().ep_join_rank_offset + self.ps.tp_size
join_effective_ep_size = get_parallel().ep_join_rank_offset + self.tp_size
dist.barrier(group=self.tp_group.cpu_group)
if self.ps.tp_rank == 0:
if self.tp_rank == 0:
register_scale_cohort(
get_parallel().ep_join_rank_offset,
join_effective_ep_size,
@@ -513,7 +513,7 @@ class ModelRunner:
join_scale_process_group()
get_context().override("elastic_ep.scale_join", ep_size=join_effective_ep_size)
global_ep_rank = self.ps.tp_rank + get_parallel().ep_join_rank_offset
global_ep_rank = self.tp_rank + get_parallel().ep_join_rank_offset
broadcast_global_expert_location_metadata(
model_config=self.model_config,
moe_ep_rank=global_ep_rank,
@@ -563,7 +563,7 @@ class ModelRunner:
def init_weight_updater(self):
self.weight_updater = WeightUpdater(
tp_rank=self.ps.tp_rank,
tp_rank=self.tp_rank,
device=self.device,
gpu_id=self.gpu_id,
model_config=self.model_config,
@@ -586,8 +586,8 @@ class ModelRunner:
def init_weight_exporter(self):
self.weight_exporter = WeightExporter(
tp_rank=self.ps.tp_rank,
tp_size=self.ps.tp_size,
tp_rank=self.tp_rank,
tp_size=self.tp_size,
gpu_id=self.gpu_id,
get_model_path=lambda: self.model_config.model_path,
get_model=lambda: self.model,
@@ -596,7 +596,7 @@ class ModelRunner:
def init_remote_instance_weight_transporter(self):
self.remote_instance_weight_transporter = RemoteInstanceWeightTransporter(
get_model=lambda: self.model,
tp_rank=self.ps.tp_rank,
tp_rank=get_parallel().tp_rank,
gpu_id=self.gpu_id,
)
@@ -610,10 +610,13 @@ class ModelRunner:
)
def init_kv_cache_configurator(self):
# The replica count that shares this KV budget is the deployment's, not
# the one a draft scope reports; the pool is allocated outside any scope.
self.kv_cache_configurator = KVCacheConfigurator(
device=self.device,
gpu_id=self.gpu_id,
ps=self.ps,
attn_dp_size=get_parallel().attn_dp_size,
pp_size=self.pp_size,
pp_group=self.pp_group,
model=self.model,
model_config=self.model_config,
@@ -647,8 +650,8 @@ class ModelRunner:
from sglang.srt.model_executor.mindspore_runner import init_ms_distributed
init_ms_distributed(
world_size=self.ps.tp_size * get_parallel().pp_size,
rank=self.ps.tp_size * get_parallel().pp_rank + self.ps.tp_rank,
world_size=self.tp_size * get_parallel().pp_size,
rank=self.tp_size * get_parallel().pp_rank + self.tp_rank,
local_rank=self.gpu_id,
port=self.dist_port,
)
@@ -718,7 +721,7 @@ class ModelRunner:
moe_ep_rank=expert_rank,
)
)
if self.ps.tp_rank == 0 and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get():
if self.tp_rank == 0 and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get():
logger.info(
"Initial expert_location_metadata:\n%s",
format_expert_location_layout(get_global_expert_location_metadata()),
@@ -779,7 +782,7 @@ class ModelRunner:
def maybe_apply_post_load_model_transforms(self):
supports_torch_tp = getattr(self.model, "supports_torch_tp", False)
if self.ps.tp_size > 1 and supports_torch_tp:
if self.tp_size > 1 and supports_torch_tp:
self.apply_torch_tp()
def maybe_init_lora_manager(self):
@@ -1015,7 +1018,7 @@ class ModelRunner:
def post_capture_elastic_ep_recover(self):
join_process_groups()
global_ep_rank = self.ps.tp_rank + get_parallel().ep_join_rank_offset
global_ep_rank = self.tp_rank + get_parallel().ep_join_rank_offset
broadcast_global_expert_location_metadata(
model_config=self.model_config,
moe_ep_rank=global_ep_rank,
@@ -1162,13 +1165,13 @@ class ModelRunner:
def check_quantized_moe_compatibility(self):
check_quantized_moe_compatibility(
model_config=self.model_config,
tp_size=self.ps.tp_size,
tp_size=self.tp_size,
moe_ep_size=get_parallel().moe_ep_size,
moe_dp_size=get_parallel().moe_dp_size,
)
def init_torch_distributed(self):
result = bootstrap.init_torch_distributed(
self.pre_model_load_memory = bootstrap.init_torch_distributed(
server_args=self.server_args,
model_config=self.model_config,
device=self.device,
@@ -1177,10 +1180,23 @@ class ModelRunner:
is_draft_worker=self.is_draft_worker,
local_omp_cpuid=self.local_omp_cpuid if self.device == "cpu" else None,
)
self.tp_group = result.tp_group
self.pp_group = result.pp_group
self.attention_tp_group = result.attention_tp_group
self.pre_model_load_memory = result.pre_model_load_memory
# Read once, here: a draft runner is constructed inside the scope that
# states its topology and used outside it, so what it holds has to be
# the placement it was built for rather than whatever the context
# answers later. Groups and widths alike -- a runner asked about its own
# shape after the scope has closed must still describe itself.
parallel = get_parallel()
self.tp_group = parallel.tp_group
self.pp_group = parallel.pp_group
self.attention_tp_group = parallel.attn_tp_group
self.tp_rank = parallel.tp_rank
self.tp_size = parallel.tp_size
self.dp_size = parallel.dp_size
self.attn_dp_size = parallel.attn_dp_size
self.pp_rank = parallel.pp_rank
self.pp_size = parallel.pp_size
self.attn_cp_rank = parallel.attn_cp_rank
self.attn_cp_size = parallel.attn_cp_size
def init_shared_mooncake_transfer_engine(self):
maybe_init_shared_mooncake_transfer_engine(gpu_id=self.gpu_id)
@@ -1204,7 +1220,7 @@ class ModelRunner:
self.load_config = build_load_config(
server_args=self.server_args,
load_format=draft_load_format,
tp_rank=self.ps.tp_rank,
tp_rank=self.tp_rank,
remote_instance_weight_transporter_engine=self.remote_instance_weight_transporter.engine,
remote_instance_weight_transporter_session_id=self.remote_instance_weight_transporter.session_id,
draft_model_idx=self.draft_model_idx,
@@ -1217,11 +1233,11 @@ class ModelRunner:
)
if self.device == "cpu":
self.model_config = adjust_config_with_unaligned_cpu_tp(
self.model_config, self.load_config, self.ps.tp_size
self.model_config, self.load_config, self.tp_size
)
maybe_trigger_remote_instance_nccl_send_group(
tp_rank=self.ps.tp_rank,
tp_rank=self.tp_rank,
load_format=draft_load_format,
)
@@ -1297,8 +1313,8 @@ class ModelRunner:
model=self.model,
spec_algorithm=self.spec_algorithm,
is_draft_worker=self.is_draft_worker,
tp_size=self.ps.tp_size,
tp_rank=self.ps.tp_rank,
tp_size=self.tp_size,
tp_rank=self.tp_rank,
pp_rank=get_parallel().pp_rank,
)
@@ -1316,7 +1332,7 @@ class ModelRunner:
if self.startup_weight_load is None:
dist_barrier_after_load(
elastic_ep_backend=get_exec().moe.elastic_ep_backend,
tp_rank=self.ps.tp_rank,
tp_rank=self.tp_rank,
is_ep_joiner=get_exec().moe.is_ep_joiner,
)
@@ -1339,7 +1355,7 @@ class ModelRunner:
self.startup_weight_load.finalize()
dist_barrier_after_load(
elastic_ep_backend=get_exec().moe.elastic_ep_backend,
tp_rank=self.ps.tp_rank,
tp_rank=self.tp_rank,
is_ep_joiner=get_exec().moe.is_ep_joiner,
)
self.startup_weight_load = None
@@ -1367,8 +1383,8 @@ class ModelRunner:
dtype=self.dtype,
server_args=self.server_args,
lora_backend=get_lora().lora_backend,
tp_size=self.ps.tp_size,
tp_rank=self.ps.tp_rank,
tp_size=self.tp_size,
tp_rank=self.tp_rank,
max_lora_rank=get_lora().max_lora_rank,
target_modules=get_lora().lora_target_modules,
lora_paths=get_lora().lora_paths,
@@ -1583,15 +1599,16 @@ class ModelRunner:
# rather than spawning additional processes, so dp_size must not be
# multiplied into the process count here (unlike regular DP, where
# dp_size * tp_size * pp_size is the true worker count).
dp_size = 1 if get_parallel().enable_dp_attention else self.ps.dp_size
parallel = get_parallel()
dp_size = 1 if parallel.enable_dp_attention else parallel.dp_size
self.local_omp_cpuid = numa_utils.init_threads_binding(
numa_index=self.gpu_id,
world_size=dp_size * self.ps.tp_size * get_parallel().pp_size,
world_size=dp_size * parallel.tp_size * parallel.pp_size,
)
def apply_torch_tp(self):
model_parallel.apply_torch_tp(
model=self.model, device=self.device, tp_size=self.ps.tp_size
model=self.model, device=self.device, tp_size=self.tp_size
)
def update_decode_attn_backend(self, stream_idx: int):
@@ -1717,7 +1734,7 @@ class ModelRunner:
if self.msprobe_debugger is not None:
rank_id = (
self.gpu_id
if self.ps.attn_dp_size is not None and self.ps.attn_dp_size > 1
if self.attn_dp_size is not None and self.attn_dp_size > 1
else None
)
self.msprobe_debugger.start(model=self.model, rank_id=rank_id)
@@ -2098,7 +2115,7 @@ class ModelRunner:
set_global_expert_location_metadata(new_metadata, allow_overwrite=True)
def _elastic_global_rank(self) -> int:
return self.ps.tp_rank + get_parallel().ep_join_rank_offset
return self.tp_rank + get_parallel().ep_join_rank_offset
def _rearm_eplb_after_elastic_scale(self) -> None:
if self.eplb_manager is None:
@@ -2120,7 +2137,7 @@ class ModelRunner:
self._rearm_eplb_after_elastic_scale()
def _report_elastic_scale_failure(self, error: str, effective_size: int) -> None:
if self.ps.tp_rank != 0 or get_exec().moe.is_ep_scale_joiner:
if self.tp_rank != 0 or get_exec().moe.is_ep_scale_joiner:
return
from sglang.srt.managers.io_struct import ElasticScaleUpdateReq
@@ -2131,7 +2148,7 @@ class ModelRunner:
)
def _elastic_scale_ready_barrier(self, target_size: int, log_tag: str) -> None:
if self.ps.tp_rank == 0:
if self.tp_rank == 0:
logger.debug(
"[Elastic EP][scale] %s entering post-scale WORLD barrier "
"(target_ep_size=%d)",
@@ -2139,7 +2156,7 @@ class ModelRunner:
target_size,
)
dist.barrier(group=dist.group.WORLD)
if self.ps.tp_rank == 0:
if self.tp_rank == 0:
logger.debug(
"[Elastic EP][scale] %s passed post-scale WORLD barrier "
"(target_ep_size=%d)",
@@ -2204,7 +2221,7 @@ class ModelRunner:
ElasticEPStateManager.commit_scale()
self._rearm_eplb_after_elastic_scale()
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
from sglang.srt.managers.io_struct import ElasticScaleUpdateReq
self._pending_elastic_scale_update = ElasticScaleUpdateReq(
@@ -2237,7 +2254,7 @@ class ModelRunner:
)
ElasticEPStateManager.fail_recovery(error)
self._report_elastic_scale_failure(error, effective_size)
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
logger.error("[Elastic EP] %s", error)
return
@@ -2263,7 +2280,7 @@ class ModelRunner:
ElasticEPStateManager.fail_scale(error)
self._reset_eplb_after_elastic_scale_failure()
self._report_elastic_scale_failure(error, effective_size)
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
logger.error("[Elastic EP] %s", error)
return
@@ -2279,7 +2296,7 @@ class ModelRunner:
ElasticEPStateManager.fail_scale(error)
self._reset_eplb_after_elastic_scale_failure()
self._report_elastic_scale_failure(error, effective_size)
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
logger.error("[Elastic EP] %s", error)
return
if not ElasticEPStateManager.begin_scale():
python/sglang/srt/utils/weight_checker.py
+25 -12
View File
@@ -8,6 +8,7 @@ import torch.distributed as dist
from pydantic import BaseModel, ConfigDict
from sglang.srt.managers.mm_utils import tensor_hash
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils.weight_checker_comparator import (
CHUNK_NUMEL,
ComparableWeight,
@@ -67,9 +68,23 @@ def _is_non_persistent_buffer_name(name: str) -> bool:
class WeightChecker:
def __init__(self, *, get_model: Callable[[], Any], ps: Any):
def __init__(self, *, get_model: Callable[[], Any]):
self._get_model = get_model
self._ps = ps
# A check is served on demand from the scheduler loop, which is outside
# the scope that describes a draft runner. The report has to name the
# runner it was built for, so the placement is read here, at
# construction, rather than asked for when the request arrives.
parallel = get_parallel()
self._placement = ParallelismInfo(
tp_rank=parallel.tp_rank,
tp_size=parallel.tp_size,
dp_rank=parallel.dp_rank if parallel.dp_rank is not None else 0,
dp_size=parallel.attn_dp_size,
pp_rank=parallel.pp_rank,
pp_size=parallel.pp_size,
rank=0,
size=1,
)
self._snapshot_tensors = None
def handle(self, action: str, allow_quant_error: bool = False) -> Optional[Dict]:
@@ -161,16 +176,14 @@ class WeightChecker:
return info.model_dump()
def _parallelism_info(self) -> ParallelismInfo:
ps = self._ps
return ParallelismInfo(
tp_rank=ps.tp_rank,
tp_size=ps.tp_size,
dp_rank=ps.dp_rank if ps.dp_rank is not None else 0,
dp_size=ps.attn_dp_size,
pp_rank=ps.pp_rank,
pp_size=ps.pp_size,
rank=dist.get_rank() if dist.is_initialized() else 0,
size=dist.get_world_size() if dist.is_initialized() else 1,
# The WORLD position is asked for now rather than frozen: unlike the
# runner's placement it is a property of the process, and an elastic
# scale-up moves it.
return self._placement.model_copy(
update={
"rank": dist.get_rank() if dist.is_initialized() else 0,
"size": dist.get_world_size() if dist.is_initialized() else 1,
}
)
def _model_state(self):