Skip to content

Roadmap: post-alpha planned work #1

Description

@MayDomine

Tracking issue for planned work after the 0.1.0-alpha open-source release. Each section lists the capability, why it matters for this architecture. Nothing here is committed; discuss scope in the comments.

Training

Rollout Router Replay: tackle MoE training–inference mismatch

The rollout/inference path and the trainer path re-route MoE tokens independently, so expert assignment drifts between the two engines and the surrogate ratio is computed against a policy the trainer never actually ran. Router replay records the routing decisions made during rollout and forces the trainer to reuse them, which removes that drift.

Dynamic CP: dynamic context parallelism for variable-length sequences

cp.py shards [B, S] through torchtitan's head-tail load balancer under a fixed cp_degree, which balances well only when sequence lengths inside a micro-batch are already uniform. RL rollouts are anything but uniform. Dynamic CP would re-partition per micro-batch — and possibly re-form the CP group — so that every rank gets a comparable amount of real (non-pad) work.

OPD

Multi-Teacher OPD (MOPD)

OPDTeacherService currently models a single frozen Teacher. MOPD lets a Student distill from several Teachers at once — different model families, different checkpoints of the same run, or a curriculum over them — which needs topology support for N teacher services, per-teacher weighting, and a merge step for the union of their score columns.

Full-Vocabulary OPD

The shipped recipe scores the Student against the Teacher over a Top-K support only, which keeps the transfer small but leaves the tail of the distribution unsupervised. Full-vocabulary OPD transfers the Teacher's complete logit distribution; the open question is bandwidth — either the Teacher column grows to vocab-size per row, or the Teacher runs the forward KL itself and only the resulting gradient proxy crosses the queue.

DataFlow

TransferQueue with RDMA backend via Mooncake

TransferQueue is the single data and control plane, so its throughput is the framework's throughput. The current path moves tensors through host memory; an RDMA backend built on Mooncake would let tensors go GPU-to-GPU and take the CPU copy and the host-bandwidth ceiling out of the critical path. This touches the vendored meshy/transferqueue/ layer.

Rollout

White-box and black-box agentic rollout

AgentLoopService / scripts/smoke/agentloop.py are the current driver, aimed at environments the framework owns end-to-end. Supporting agents generally means two shapes: white-box (the environment runs in-process, token-level access, per-step log-probs available to the trainer) and black-box (a remote HTTP API returns only text, no token ids and no log-probs). Black-box rollout needs a way to score a trajectory after the fact rather than during it, and an import path for externally produced trajectories.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions