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

Compare commits

base: zacheryasc:6c11c4c2b4534bc2f7e3e731200a046dc2f878ac
Branches Tags
zacheryasc:master
zacheryasc:development
zacheryasc:backup/pre-rewrite-msg
..
compare: zacheryasc:06478bb6a2d2d73a7d2ef8c3acbbe28c0d27a633
Branches Tags
zacheryasc:development
zacheryasc:master
zacheryasc:backup/pre-rewrite-msg

No commits in common. "6c11c4c2b4534bc2f7e3e731200a046dc2f878ac" and "06478bb6a2d2d73a7d2ef8c3acbbe28c0d27a633" have entirely different histories.
6c11c4c2b4 ... 06478bb6a2

3 changed files with 139 additions and 127 deletions
Show stats Expand all files Collapse all files

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

@ -395,48 +395,30 @@ class PipelineStageTinygradModel:
first_stage: bool, first_stage: bool,
final_stage: bool, final_stage: bool,
nn_mod: Any, nn_mod: Any,
config_cls: Any | None, config_cls: Any,
block_cls: Any, block_cls: Any,
) -> None: ) -> None:
if config_cls is None: block_config = config_cls(
self.blk = [ num_blocks=block_count,
block_cls( dim=dim,
dim, hidden_dim=hidden_dim,
hidden_dim, n_heads=n_heads,
n_heads, n_kv_heads=n_kv_heads,
n_kv_heads, norm_eps=norm_eps,
norm_eps, vocab_size=vocab_size,
head_dim, head_dim=head_dim,
rope_theta, rope_theta=rope_theta,
max_context, rope_dim=rope_dim,
qk_norm, v_head_dim=v_head_dim,
num_experts, max_context=max_context,
num_experts_per_tok, qk_norm=qk_norm,
) num_experts=num_experts,
for _ in range(block_count) num_experts_per_tok=num_experts_per_tok,
] norm_topk_prob=norm_topk_prob,
else: qkv_bias=qkv_bias,
block_config = config_cls( expert_bias=expert_bias,
num_blocks=block_count, )
dim=dim, self.blk = [block_cls(block_config) for _ in range(block_count)]
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
@ -497,22 +479,10 @@ 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
try: kv, state_dict = gguf_load(path)
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"]
@ -568,8 +538,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=config_cls, config_cls=TransformerConfig,
block_cls=block_cls, block_cls=TransformerBlock,
) )
stage_state = remap_stage_state_dict( stage_state = remap_stage_state_dict(
state_dict, state_dict,
@ -609,11 +579,12 @@ 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 is not None and not exc.name.startswith("tinygrad.llm"): if exc.name != "tinygrad.llm":
raise raise
from tinygrad.apps.llm import SimpleTokenizer from tinygrad.apps.llm import SimpleTokenizer, Transformer
llm_backend = "tinygrad.apps.llm" llm_backend = "tinygrad.apps.llm"
@ -631,12 +602,28 @@ 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,
) )
model, kv = load_pipeline_stage_model( if Transformer is None:
path, model, kv = load_pipeline_stage_model(
max_context=max_context, path,
layer_start=layer_start, max_context=max_context,
layer_end_exclusive=layer_end_exclusive, 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
control( control(
type="PipelineStageFromGgufReady", type="PipelineStageFromGgufReady",
model_id=model_id, model_id=model_id,

21
crates/mvp-system/src/orchestrator_app.rs
View file

@ -4660,7 +4660,6 @@ fn stop_requested(stop_rx: &mpsc::Receiver<()>) -> bool {
fn spawn_stop_listener() -> mpsc::Receiver<()> { fn spawn_stop_listener() -> mpsc::Receiver<()> {
let (tx, rx) = mpsc::channel(); let (tx, rx) = mpsc::channel();
let stdin_tx = tx.clone();
thread::spawn(move || { thread::spawn(move || {
let stdin = std::io::stdin(); let stdin = std::io::stdin();
for line in stdin.lock().lines().map_while(Result::ok) { for line in stdin.lock().lines().map_while(Result::ok) {
@ -4669,29 +4668,11 @@ fn spawn_stop_listener() -> mpsc::Receiver<()> {
|| trimmed.eq_ignore_ascii_case("shutdown") || trimmed.eq_ignore_ascii_case("shutdown")
|| trimmed.eq_ignore_ascii_case("quit") || trimmed.eq_ignore_ascii_case("quit")
{ {
let _ = stdin_tx.send(()); let _ = tx.send(());
break; break;
} }
} }
}); });
#[cfg(target_os = "linux")]
{
thread::spawn(move || {
let Ok(mut signals) = signal_hook::iterator::Signals::new([
signal_hook::consts::signal::SIGINT,
signal_hook::consts::signal::SIGTERM,
]) else {
return;
};
if signals.forever().next().is_some() {
let _ = tx.send(());
}
});
}
#[cfg(not(target_os = "linux"))]
{
drop(tx);
}
rx rx
} }

128
xtask/src/main.rs
View file

@ -86,6 +86,12 @@ 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) {
@ -1238,7 +1244,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_network_facts(&facts)?; require_multinode_docker_relay_facts(&facts)?;
} }
Ok(events) Ok(events)
} }
@ -1289,14 +1295,12 @@ fn build_benchmark_report(
"ready", "ready",
)?; )?;
} }
if scenario != MvpChatCheckScenario::MultinodeDocker { facts.require_span(
facts.require_span( "mvp.chat.runtime",
"mvp.chat.runtime", "ChatProgress",
"ChatProgress", "ensure_worker_binary",
"ensure_worker_binary", "ready",
"ready", )?;
)?;
}
facts.require_span( facts.require_span(
"mvp.orch.bootstrap", "mvp.orch.bootstrap",
"OrchBootstrap", "OrchBootstrap",
@ -1590,9 +1594,10 @@ 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,
docker_node_spec_worker_count: Option<u64>, relay_masked_orchestrator_ready: bool,
docker_worker_iroh_ready: BTreeSet<u64>, relay_masked_node_spec_worker_count: Option<u64>,
docker_worker_coordinator_join: BTreeSet<u64>, relay_masked_worker_iroh_ready: 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,
@ -1638,23 +1643,31 @@ 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 let Some(node_id) = event_node_id(event) { if detail_relay_only_advertisement(event)
facts.docker_worker_iroh_ready.insert(node_id); && let Some(node_id) = event_node_id(event)
{
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, "provider") == Some("docker") { if detail_str(event, "endpoint_addr_mask") == Some("relay-only")
facts.docker_node_spec_worker_count = detail_u64(event, "worker_count"); && detail_str(event, "relay_mode") == Some("default")
{
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_u64(event, "direct_addr_count").is_some_and(|count| count > 0) if detail_bool(event, "has_relay") == Some(true)
&& 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.docker_worker_coordinator_join.insert(node_id); facts.relay_masked_worker_coordinator_join.insert(node_id);
} }
} }
(_, Some("NodeEvent"), Some("worker_initialize"), Some("ready")) => { (_, Some("NodeEvent"), Some("worker_initialize"), Some("ready")) => {
@ -1865,20 +1878,24 @@ fn require_gpu_dump_log_facts(facts: &DumpLogFacts) -> Result<(), String> {
Ok(()) Ok(())
} }
fn require_multinode_docker_network_facts(facts: &DumpLogFacts) -> Result<(), String> { fn require_multinode_docker_relay_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
.docker_node_spec_worker_count .relay_masked_node_spec_worker_count
.is_some_and(|count| count >= 2), .is_some_and(|count| count >= 2),
"Docker node_spec with multiple workers", "relay-masked Docker node_spec with multiple workers",
)?; )?;
require_dump_log_fact( require_dump_log_fact(
facts.docker_worker_iroh_ready.len() >= 2, facts.relay_masked_worker_iroh_ready.len() >= 2,
"Docker worker iroh_driver ready for multiple nodes", "relay-masked worker iroh_driver ready for multiple nodes",
)?; )?;
require_dump_log_fact( require_dump_log_fact(
facts.docker_worker_coordinator_join.len() >= 2, facts.relay_masked_worker_coordinator_join.len() >= 2,
"Docker worker direct coordinator_join for multiple nodes", "relay-masked worker coordinator_join for multiple nodes",
) )
} }
@ -1944,6 +1961,13 @@ 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")
@ -1958,6 +1982,12 @@ 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")
@ -2041,6 +2071,10 @@ 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",
@ -2792,13 +2826,22 @@ mod tests {
} }
#[test] #[test]
fn benchmark_observability_multinode_docker_dump_facts_require_direct_network_events() { fn benchmark_observability_multinode_docker_dump_facts_require_relay_masked_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":"node_spec","status":"ready","run_id":9,"node_id":1,"detail":{"endpoint_addr_mask":"full","provider":"docker","worker_count":2}}), 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}}),
"mvp-orchestrator", "mvp-orchestrator",
1_071, 1_071,
71, 71,
@ -2807,7 +2850,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":"full","has_relay":false,"direct_addr_count":3}}), 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}}),
"mvp-worker-node", "mvp-worker-node",
1_072, 1_072,
72, 72,
@ -2816,7 +2859,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":"full","has_relay":false,"direct_addr_count":3}}), 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}}),
"mvp-worker-node", "mvp-worker-node",
1_073, 1_073,
73, 73,
@ -2825,7 +2868,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":false,"direct_addr_count":4}}), 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}}),
"mvp-worker-node", "mvp-worker-node",
1_074, 1_074,
74, 74,
@ -2834,35 +2877,36 @@ 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":false,"direct_addr_count":4}}), 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}}),
"mvp-worker-node", "mvp-worker-node",
1_075, 1_075,
75, 75,
), ),
), ),
]); ]);
let path = write_synthetic_event_dump("multinode-docker-direct-network", events); let path = write_synthetic_event_dump("multinode-docker-relay-mask", events);
assert_dump_log_facts(&path, MvpChatCheckScenario::MultinodeDocker) assert_dump_log_facts(&path, MvpChatCheckScenario::MultinodeDocker)
.expect("direct-network multinode Docker facts pass"); .expect("relay-masked multinode Docker facts pass");
let _ = fs::remove_file(path); let _ = fs::remove_file(path);
} }
#[test] #[test]
fn benchmark_observability_multinode_docker_requires_direct_network_workers() { fn benchmark_observability_multinode_docker_requires_relay_masked_workers() {
let mut valid = DumpLogFacts { let mut valid = DumpLogFacts {
docker_node_spec_worker_count: Some(2), relay_masked_orchestrator_ready: true,
relay_masked_node_spec_worker_count: Some(2),
..DumpLogFacts::default() ..DumpLogFacts::default()
}; };
valid.docker_worker_iroh_ready.extend([2, 3]); valid.relay_masked_worker_iroh_ready.extend([2, 3]);
valid.docker_worker_coordinator_join.extend([2, 3]); valid.relay_masked_worker_coordinator_join.extend([2, 3]);
require_multinode_docker_network_facts(&valid).expect("direct-network Docker facts pass"); require_multinode_docker_relay_facts(&valid).expect("relay-masked Docker facts pass");
let mut missing_worker = valid; let mut missing_worker = valid;
missing_worker.docker_worker_iroh_ready.remove(&3); missing_worker.relay_masked_worker_iroh_ready.remove(&3);
let error = require_multinode_docker_network_facts(&missing_worker) let error = require_multinode_docker_relay_facts(&missing_worker)
.expect_err("single direct-network worker should fail"); .expect_err("single relay-masked worker should fail");
assert!( assert!(
error.contains("Docker worker iroh_driver ready for multiple nodes"), error.contains("relay-masked worker iroh_driver ready for multiple nodes"),
"unexpected error: {error}" "unexpected error: {error}"
); );
} }