zacheryasc/swactor
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: working multinode pipeline parallel prompt loop running locally

Browse source 
Get the multinode pipeline-parallel prompt loop actually running locally by supporting both tinygrad LLM backends and switching the Docker scenario to direct addressing.

- tinygrad_worker.py (load_pipeline_stage_model): try the modern tinygrad.llm gguf/model modules first and fall back to the legacy tinygrad.apps.llm TransformerBlock on ModuleNotFoundError, with PipelineStageTinygradModel constructing blocks positionally when no TransformerConfig exists
- tinygrad_worker.py (load_weights): drop the Transformer.from_gguf whole-model branch and its TinygradAppsLlmPartialStageUnsupported fatal, so partial pipeline stages build through load_pipeline_stage_model on either backend
- xtask (MultinodeDocker): stop forcing --relay-mode default --endpoint-addr-mask relay-only, so the scenario runs over the Docker network with full/direct addresses
- xtask (dump-log facts): rename relay_masked_* facts to docker_*, assert multiple workers join the coordinator via direct addresses (direct_addr_count > 0), and relax the benchmark report to skip the ensure_worker_binary span for the Docker scenario

Signed-off-by: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-23 15:34:57 +04:00
parent 42bf9cceff
commit fdc3c9663f
2 changed files with 107 additions and 138 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

117
apps/mvp-node/tinygrad_worker.py
View file

@ -395,30 +395,48 @@ class PipelineStageTinygradModel:
first_stage: bool,
final_stage: bool,
nn_mod: Any,
config_cls: Any,
config_cls: Any | None,
block_cls: Any,
) -> None:
block_config = config_cls(
num_blocks=block_count,
dim=dim,
hidden_dim=hidden_dim,
n_heads=n_heads,
n_kv_heads=n_kv_heads,
norm_eps=norm_eps,
vocab_size=vocab_size,
head_dim=head_dim,
rope_theta=rope_theta,
rope_dim=rope_dim,
v_head_dim=v_head_dim,
max_context=max_context,
qk_norm=qk_norm,
num_experts=num_experts,
num_experts_per_tok=num_experts_per_tok,
norm_topk_prob=norm_topk_prob,
qkv_bias=qkv_bias,
expert_bias=expert_bias,
)
self.blk = [block_cls(block_config) for _ in range(block_count)]
if config_cls is None:
self.blk = [
block_cls(
dim,
hidden_dim,
n_heads,
n_kv_heads,
norm_eps,
head_dim,
rope_theta,
max_context,
qk_norm,
num_experts,
num_experts_per_tok,
)
for _ in range(block_count)
]
else:
block_config = config_cls(
num_blocks=block_count,
dim=dim,
hidden_dim=hidden_dim,
n_heads=n_heads,
n_kv_heads=n_kv_heads,
norm_eps=norm_eps,
vocab_size=vocab_size,
head_dim=head_dim,
rope_theta=rope_theta,
rope_dim=rope_dim,
v_head_dim=v_head_dim,
max_context=max_context,
qk_norm=qk_norm,
num_experts=num_experts,
num_experts_per_tok=num_experts_per_tok,
norm_topk_prob=norm_topk_prob,
qkv_bias=qkv_bias,
expert_bias=expert_bias,
)
self.blk = [block_cls(block_config) for _ in range(block_count)]
self.max_context = max_context
self.hidden_dim = dim
self.first_stage = first_stage
@ -479,10 +497,22 @@ def load_pipeline_stage_model(
) -> tuple[PipelineStageTinygradModel, dict[str, Any]]:
TensorCls = require_tinygrad()
from tinygrad import nn
from tinygrad.llm.gguf import gguf_load
from tinygrad.llm.model import TransformerBlock, TransformerConfig
kv, state_dict = gguf_load(path)
try:
from tinygrad.llm.gguf import gguf_load
from tinygrad.llm.model import TransformerBlock, TransformerConfig
kv, state_dict = gguf_load(path)
block_cls = TransformerBlock
config_cls = TransformerConfig
except ModuleNotFoundError as exc:
if exc.name is not None and not exc.name.startswith("tinygrad.llm"):
raise
from tinygrad.apps.llm import TransformerBlock
kv, state_dict = nn.state.gguf_load(TensorCls(path).to(None))
block_cls = TransformerBlock
config_cls = None
state_dict = {key: value.cast("float16") if env_flag("HALF", True) else value for key, value in state_dict.items()}
if "output.weight" not in state_dict and "token_embd.weight" in state_dict:
state_dict["output.weight"] = state_dict["token_embd.weight"]
@ -538,8 +568,8 @@ def load_pipeline_stage_model(
first_stage=first_stage,
final_stage=final_stage,
nn_mod=nn,
config_cls=TransformerConfig,
block_cls=TransformerBlock,
config_cls=config_cls,
block_cls=block_cls,
)
stage_state = remap_stage_state_dict(
state_dict,
@ -579,12 +609,11 @@ def load_weights(cmd: dict[str, Any]) -> None:
try:
from tinygrad.llm.cli import SimpleTokenizer
Transformer = None
llm_backend = "tinygrad.llm"
except ModuleNotFoundError as exc:
if exc.name != "tinygrad.llm":
if exc.name is not None and not exc.name.startswith("tinygrad.llm"):
raise
from tinygrad.apps.llm import SimpleTokenizer, Transformer
from tinygrad.apps.llm import SimpleTokenizer
llm_backend = "tinygrad.apps.llm"
@ -602,28 +631,12 @@ def load_weights(cmd: dict[str, Any]) -> None:
requested_device=os.environ.get("DEV"),
llm_backend=llm_backend,
)
if Transformer is None:
model, kv = load_pipeline_stage_model(
path,
max_context=max_context,
layer_start=layer_start,
layer_end_exclusive=layer_end_exclusive,
)
else:
TensorCls = require_tinygrad()
model, kv = Transformer.from_gguf(TensorCls(path), max_context=max_context, realize=True)
total_layers = int(kv[f"{kv['general.architecture']}.block_count"]) - int(
kv.get(f"{kv['general.architecture']}.nextn_predict_layers", 0)
)
if layer_start != 0 or layer_end_exclusive < total_layers:
fatal(
"TinygradAppsLlmPartialStageUnsupported",
layer_start=layer_start,
layer_end_exclusive=layer_end_exclusive,
total_layers=total_layers,
)
model.first_stage = True
model.final_stage = True
model, kv = load_pipeline_stage_model(
path,
max_context=max_context,
layer_start=layer_start,
layer_end_exclusive=layer_end_exclusive,
)
control(
type="PipelineStageFromGgufReady",
model_id=model_id,

128
xtask/src/main.rs
View file

@ -86,12 +86,6 @@ impl MvpChatCheckScenario {
}
Self::MultinodeDocker => {
args.push("--docker".to_owned());
args.extend([
"--relay-mode".to_owned(),
"default".to_owned(),
"--endpoint-addr-mask".to_owned(),
"relay-only".to_owned(),
]);
}
}
if matches!(self, Self::Multinode | Self::MultinodeDocker) {
@ -1244,7 +1238,7 @@ fn assert_dump_log_facts(
require_gpu_dump_log_facts(&facts)?;
}
if scenario == MvpChatCheckScenario::MultinodeDocker {
require_multinode_docker_relay_facts(&facts)?;
require_multinode_docker_network_facts(&facts)?;
}
Ok(events)
}
@ -1295,12 +1289,14 @@ fn build_benchmark_report(
"ready",
)?;
}
facts.require_span(
"mvp.chat.runtime",
"ChatProgress",
"ensure_worker_binary",
"ready",
)?;
if scenario != MvpChatCheckScenario::MultinodeDocker {
facts.require_span(
"mvp.chat.runtime",
"ChatProgress",
"ensure_worker_binary",
"ready",
)?;
}
facts.require_span(
"mvp.orch.bootstrap",
"OrchBootstrap",
@ -1594,10 +1590,9 @@ struct DumpLogFacts {
gpu_pipeline_tokenizer_decode_ready: BTreeSet<u64>,
gpu_pipeline_tokens_decoded: BTreeSet<u64>,
gpu_pipeline_real_worker_step_seen: bool,
relay_masked_orchestrator_ready: bool,
relay_masked_node_spec_worker_count: Option<u64>,
relay_masked_worker_iroh_ready: BTreeSet<u64>,
relay_masked_worker_coordinator_join: BTreeSet<u64>,
docker_node_spec_worker_count: Option<u64>,
docker_worker_iroh_ready: BTreeSet<u64>,
docker_worker_coordinator_join: BTreeSet<u64>,
chat_config_ready: bool,
prepare_runtime_ready: bool,
prompt_rpc_ready: bool,
@ -1643,31 +1638,23 @@ fn record_dump_log_event(
}
(_, Some("OrchBootstrap"), Some("iroh_driver"), Some("ready")) => {
facts.orch_iroh_driver_ready = true;
if detail_relay_only_advertisement(event) {
facts.relay_masked_orchestrator_ready = true;
}
}
(_, Some("NodeEvent"), Some("iroh_driver"), Some("ready")) => {
facts.node_iroh_driver_ready = true;
if detail_relay_only_advertisement(event)
&& let Some(node_id) = event_node_id(event)
{
facts.relay_masked_worker_iroh_ready.insert(node_id);
if let Some(node_id) = event_node_id(event) {
facts.docker_worker_iroh_ready.insert(node_id);
}
}
(_, Some("OrchBootstrap"), Some("node_spec"), Some("ready")) => {
if detail_str(event, "endpoint_addr_mask") == Some("relay-only")
&& detail_str(event, "relay_mode") == Some("default")
{
facts.relay_masked_node_spec_worker_count = detail_u64(event, "worker_count");
if detail_str(event, "provider") == Some("docker") {
facts.docker_node_spec_worker_count = detail_u64(event, "worker_count");
}
}
(_, Some("NodeEvent"), Some("coordinator_join"), Some("started")) => {
if detail_bool(event, "has_relay") == Some(true)
&& detail_u64(event, "direct_addr_count") == Some(0)
if detail_u64(event, "direct_addr_count").is_some_and(|count| count > 0)
&& let Some(node_id) = event_node_id(event)
{
facts.relay_masked_worker_coordinator_join.insert(node_id);
facts.docker_worker_coordinator_join.insert(node_id);
}
}
(_, Some("NodeEvent"), Some("worker_initialize"), Some("ready")) => {
@ -1878,24 +1865,20 @@ fn require_gpu_dump_log_facts(facts: &DumpLogFacts) -> Result<(), String> {
Ok(())
}
fn require_multinode_docker_relay_facts(facts: &DumpLogFacts) -> Result<(), String> {
require_dump_log_fact(
facts.relay_masked_orchestrator_ready,
"relay-masked orchestrator endpoint",
)?;
fn require_multinode_docker_network_facts(facts: &DumpLogFacts) -> Result<(), String> {
require_dump_log_fact(
facts
.relay_masked_node_spec_worker_count
.docker_node_spec_worker_count
.is_some_and(|count| count >= 2),
"relay-masked Docker node_spec with multiple workers",
"Docker node_spec with multiple workers",
)?;
require_dump_log_fact(
facts.relay_masked_worker_iroh_ready.len() >= 2,
"relay-masked worker iroh_driver ready for multiple nodes",
facts.docker_worker_iroh_ready.len() >= 2,
"Docker worker iroh_driver ready for multiple nodes",
)?;
require_dump_log_fact(
facts.relay_masked_worker_coordinator_join.len() >= 2,
"relay-masked worker coordinator_join for multiple nodes",
facts.docker_worker_coordinator_join.len() >= 2,
"Docker worker direct coordinator_join for multiple nodes",
)
}
@ -1961,13 +1944,6 @@ fn event_node_id(event: &Value) -> Option<u64> {
event.get("node_id").and_then(Value::as_u64)
}
fn detail_bool(event: &Value, key: &str) -> Option<bool> {
event
.get("detail")
.and_then(|detail| detail.get(key))
.and_then(Value::as_bool)
}
fn detail_u64(event: &Value, key: &str) -> Option<u64> {
event
.get("detail")
@ -1982,12 +1958,6 @@ fn detail_str<'a>(event: &'a Value, key: &str) -> Option<&'a str> {
.and_then(Value::as_str)
}
fn detail_relay_only_advertisement(event: &Value) -> bool {
detail_str(event, "endpoint_addr_mask") == Some("relay-only")
&& detail_bool(event, "has_relay") == Some(true)
&& detail_u64(event, "direct_addr_count") == Some(0)
}
fn dump_log_request_id(event: &Value) -> Option<u64> {
event
.get("detail")
@ -2071,10 +2041,6 @@ mod tests {
multinode_docker.mvp_chat_args(42, dump_log),
strings(&[
"--docker",
"--relay-mode",
"default",
"--endpoint-addr-mask",
"relay-only",
"--pipeline-stages",
"2",
"--cached-model",
@ -2826,22 +2792,13 @@ mod tests {
}
#[test]
fn benchmark_observability_multinode_docker_dump_facts_require_relay_masked_events() {
fn benchmark_observability_multinode_docker_dump_facts_require_direct_network_events() {
let mut events = dump_log_fact_events(false, false);
events.extend([
(
"mvp.orch.bootstrap",
stamped(
json!({"type":"OrchBootstrap","phase":"iroh_driver","status":"ready","run_id":9,"node_id":1,"detail":{"endpoint_addr_mask":"relay-only","relay_mode":"Default","has_relay":true,"direct_addr_count":0}}),
"mvp-orchestrator",
1_070,
70,
),
),
(
"mvp.orch.bootstrap",
stamped(
json!({"type":"OrchBootstrap","phase":"node_spec","status":"ready","run_id":9,"node_id":1,"detail":{"endpoint_addr_mask":"relay-only","relay_mode":"default","worker_count":2}}),
json!({"type":"OrchBootstrap","phase":"node_spec","status":"ready","run_id":9,"node_id":1,"detail":{"endpoint_addr_mask":"full","provider":"docker","worker_count":2}}),
"mvp-orchestrator",
1_071,
71,
@ -2850,7 +2807,7 @@ mod tests {
(
"mvp.node.bootstrap",
stamped(
json!({"type":"NodeEvent","phase":"iroh_driver","status":"ready","run_id":9,"node_id":2,"stage_index":0,"detail":{"endpoint_addr_mask":"relay-only","has_relay":true,"direct_addr_count":0}}),
json!({"type":"NodeEvent","phase":"iroh_driver","status":"ready","run_id":9,"node_id":2,"stage_index":0,"detail":{"endpoint_addr_mask":"full","has_relay":false,"direct_addr_count":3}}),
"mvp-worker-node",
1_072,
72,
@ -2859,7 +2816,7 @@ mod tests {
(
"mvp.node.bootstrap",
stamped(
json!({"type":"NodeEvent","phase":"iroh_driver","status":"ready","run_id":9,"node_id":3,"stage_index":1,"detail":{"endpoint_addr_mask":"relay-only","has_relay":true,"direct_addr_count":0}}),
json!({"type":"NodeEvent","phase":"iroh_driver","status":"ready","run_id":9,"node_id":3,"stage_index":1,"detail":{"endpoint_addr_mask":"full","has_relay":false,"direct_addr_count":3}}),
"mvp-worker-node",
1_073,
73,
@ -2868,7 +2825,7 @@ mod tests {
(
"mvp.node.bootstrap",
stamped(
json!({"type":"NodeEvent","phase":"coordinator_join","status":"started","run_id":9,"node_id":2,"stage_index":0,"detail":{"has_relay":true,"direct_addr_count":0}}),
json!({"type":"NodeEvent","phase":"coordinator_join","status":"started","run_id":9,"node_id":2,"stage_index":0,"detail":{"has_relay":false,"direct_addr_count":4}}),
"mvp-worker-node",
1_074,
74,
@ -2877,36 +2834,35 @@ mod tests {
(
"mvp.node.bootstrap",
stamped(
json!({"type":"NodeEvent","phase":"coordinator_join","status":"started","run_id":9,"node_id":3,"stage_index":1,"detail":{"has_relay":true,"direct_addr_count":0}}),
json!({"type":"NodeEvent","phase":"coordinator_join","status":"started","run_id":9,"node_id":3,"stage_index":1,"detail":{"has_relay":false,"direct_addr_count":4}}),
"mvp-worker-node",
1_075,
75,
),
),
]);
let path = write_synthetic_event_dump("multinode-docker-relay-mask", events);
let path = write_synthetic_event_dump("multinode-docker-direct-network", events);
assert_dump_log_facts(&path, MvpChatCheckScenario::MultinodeDocker)
.expect("relay-masked multinode Docker facts pass");
.expect("direct-network multinode Docker facts pass");
let _ = fs::remove_file(path);
}
#[test]
fn benchmark_observability_multinode_docker_requires_relay_masked_workers() {
fn benchmark_observability_multinode_docker_requires_direct_network_workers() {
let mut valid = DumpLogFacts {
relay_masked_orchestrator_ready: true,
relay_masked_node_spec_worker_count: Some(2),
docker_node_spec_worker_count: Some(2),
..DumpLogFacts::default()
};
valid.relay_masked_worker_iroh_ready.extend([2, 3]);
valid.relay_masked_worker_coordinator_join.extend([2, 3]);
require_multinode_docker_relay_facts(&valid).expect("relay-masked Docker facts pass");
valid.docker_worker_iroh_ready.extend([2, 3]);
valid.docker_worker_coordinator_join.extend([2, 3]);
require_multinode_docker_network_facts(&valid).expect("direct-network Docker facts pass");
let mut missing_worker = valid;
missing_worker.relay_masked_worker_iroh_ready.remove(&3);
let error = require_multinode_docker_relay_facts(&missing_worker)
.expect_err("single relay-masked worker should fail");
missing_worker.docker_worker_iroh_ready.remove(&3);
let error = require_multinode_docker_network_facts(&missing_worker)
.expect_err("single direct-network worker should fail");
assert!(
error.contains("relay-masked worker iroh_driver ready for multiple nodes"),
error.contains("Docker worker iroh_driver ready for multiple nodes"),
"unexpected error: {error}"
);
}