dockyard_rl.models.policy.workers.dtensor_policy_worker_v2

MoE-capable DTensor policy worker (v2).

Extends DTensorPolicyWorkerImpl with native expert parallelism. It overrides only the two seams the v1 refactor exposed:

  • _build_device_mesh — for the dense, non-HSDP case (ep=1, dp_replicate=1) it defers to v1 verbatim (byte-identical bring-up). When expert parallelism or HSDP is requested it builds one flat world mesh and derives both the dense (dp_replicate, dp_shard, cp, tp) handles v1 consumes and the sparse (dp_replicate, efsdp, ep) mesh the routed experts shard over.

  • _parallelize — applies the B.3 expert-sharding policy to every GroupedExperts submodule (over the sparse mesh) before the dense FSDP/TP plan covers the remaining (attention / norm / router) parameters.

The MoE mesh construction and expert sharding are device-bound (live process group + CUDA grouped-GEMM); they are exercised at bring-up and tracked in hardware-deferred-validation.md (HV-2/HV-9). The pure decisions they rely on (mesh topology math, placement policy, dispatch) are unit-tested no-GPU.

Module Contents

Classes

DTensorPolicyWorkerV2Impl

DTensor policy worker with native MoE expert parallelism.

DTensorPolicyWorkerV2

API

class dockyard_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2Impl

Bases: dockyard_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorkerImpl

DTensor policy worker with native MoE expert parallelism.

moe_dims: dockyard_rl.models.dtensor.moe.MoEParallelDims

None

moe_mesh: Optional[torch.distributed.device_mesh.DeviceMesh]

None

prepare_refit_info(state_dict_info: Optional[dict[str, Any]] = None) → Optional[dict[str, Any]]

Refit metadata with fused-expert entries expanded per-expert.

Mirrors _iter_refit_state_dict on the metadata side so the inference backend’s name->(shape, dtype) map matches the streamed tensors, dtype included – both read refit_tensor_dtype. Expanded per-expert tensors are unbound views of the fused param, so they inherit its dtype.

class dockyard_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2

Bases: dockyard_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2Impl