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, first_stage: bool,
final_stage: bool, final_stage: bool,
nn_mod: Any, nn_mod: Any,
config_cls: Any, config_cls: Any | None,
block_cls: Any, block_cls: Any,
) -> None: ) -> None:
block_config = config_cls( if config_cls is None:
num_blocks=block_count, self.blk = [
dim=dim, block_cls(
hidden_dim=hidden_dim, dim,
n_heads=n_heads, hidden_dim,
n_kv_heads=n_kv_heads, n_heads,
norm_eps=norm_eps, n_kv_heads,
vocab_size=vocab_size, norm_eps,
head_dim=head_dim, head_dim,
rope_theta=rope_theta, rope_theta,
rope_dim=rope_dim, max_context,
v_head_dim=v_head_dim, qk_norm,
max_context=max_context, num_experts,
qk_norm=qk_norm, num_experts_per_tok,
num_experts=num_experts, )
num_experts_per_tok=num_experts_per_tok, for _ in range(block_count)
norm_topk_prob=norm_topk_prob, ]
qkv_bias=qkv_bias, else:
expert_bias=expert_bias, block_config = config_cls(
) num_blocks=block_count,
self.blk = [block_cls(block_config) for _ in range(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.max_context = max_context
self.hidden_dim = dim self.hidden_dim = dim
self.first_stage = first_stage self.first_stage = first_stage
@ -479,10 +497,22 @@ def load_pipeline_stage_model(
) -> tuple[PipelineStageTinygradModel, dict[str, Any]]: ) -> tuple[PipelineStageTinygradModel, dict[str, Any]]:
TensorCls = require_tinygrad() TensorCls = require_tinygrad()
from tinygrad import nn 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()} 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: if "output.weight" not in state_dict and "token_embd.weight" in state_dict:
state_dict["output.weight"] = state_dict["token_embd.weight"] state_dict["output.weight"] = state_dict["token_embd.weight"]
@ -538,8 +568,8 @@ def load_pipeline_stage_model(
first_stage=first_stage, first_stage=first_stage,
final_stage=final_stage, final_stage=final_stage,
nn_mod=nn, nn_mod=nn,
config_cls=TransformerConfig, config_cls=config_cls,
block_cls=TransformerBlock, block_cls=block_cls,
) )
stage_state = remap_stage_state_dict( stage_state = remap_stage_state_dict(
state_dict, state_dict,
@ -579,12 +609,11 @@ def load_weights(cmd: dict[str, Any]) -> None:
try: try:
from tinygrad.llm.cli import SimpleTokenizer from tinygrad.llm.cli import SimpleTokenizer
Transformer = None
llm_backend = "tinygrad.llm" llm_backend = "tinygrad.llm"
except ModuleNotFoundError as exc: except ModuleNotFoundError as exc:
if exc.name != "tinygrad.llm": if exc.name is not None and not exc.name.startswith("tinygrad.llm"):
raise raise
from tinygrad.apps.llm import SimpleTokenizer, Transformer from tinygrad.apps.llm import SimpleTokenizer
llm_backend = "tinygrad.apps.llm" llm_backend = "tinygrad.apps.llm"
@ -602,28 +631,12 @@ def load_weights(cmd: dict[str, Any]) -> None:
requested_device=os.environ.get("DEV"), requested_device=os.environ.get("DEV"),
llm_backend=llm_backend, llm_backend=llm_backend,
) )
if Transformer is None: model, kv = load_pipeline_stage_model(
model, kv = load_pipeline_stage_model( path,
path, max_context=max_context,
max_context=max_context, layer_start=layer_start,
layer_start=layer_start, layer_end_exclusive=layer_end_exclusive,
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
control( control(
type="PipelineStageFromGgufReady", type="PipelineStageFromGgufReady",
model_id=model_id, model_id=model_id,

128
xtask/src/main.rs
View file

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