Export Uzu traces as safetensors#450
Conversation
ry2009
force-pushed
the
ryan/trace-safetensors-export
branch
3 times, most recently
from
0aa5d08 to
558d57c
Compare
ry2009
force-pushed
the
ryan/trace-safetensors-export
branch
2 times, most recently
from
5a7ee19 to
ff50b1f
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ff50b1fa64
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| ]; | ||
| let request = json!({ | ||
| "add_generation_prompt": true, | ||
| "messages": [{"role": token_codec_config.user_role_name, "content": message}], |
There was a problem hiding this comment.
When ChatCodecConfig::default_system_prompt is set, InputProcessorDefault::process inserts a system message before rendering, but the exported request metadata records only the user message. In that configuration, consumers that replay the saved request with the saved template will render a different prompt and token sequence than the exported rendered_request/token tensors, making the tokenization trace internally inconsistent. Include the injected system message (or build metadata from the same resolved message list used for rendering).
Useful? React with 👍 / 👎.
| name: name.into(), | ||
| shape, | ||
| data_type: array.data_type(), | ||
| data: array.as_bytes().into(), |
There was a problem hiding this comment.
For large trace exports, this copies every produced activation tensor into a boxed byte buffer and the caller keeps all of those copies in tensors until write_trace_file finishes. Since the original ActivationTrace arrays are already resident, long prompts or larger models can roughly double the trace payload in memory and fail before the file is written; the writer can compute offsets from borrowed byte slices and stream the existing arrays instead.
Useful? React with 👍 / 👎.
ry2009
force-pushed
the
ryan/trace-safetensors-export
branch
from
ff50b1f to
8771473
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 87714733b3
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| ]; | ||
| Self::push_activation_trace_tensors( | ||
| &mut tensors, | ||
| &traces, | ||
| &ctx.model_config.decoder_config.transformer_config.layer_configs, |
There was a problem hiding this comment.
When exporting decoder traces for models with attention, SSM, short-conv, or DeltaNet layers, this output is built only from token tensors plus ActivationTrace, but ActivationTrace contains activations/logits and not the updated cache state. The existing trace schema includes activation_trace.layer_results.N.updated_state.{keys,values,conv_state,ssm_state} (the previous validator read these from the cache layers), so exported safetensors for those models will be missing required tensors for downstream trace comparison/replay. Add the relevant cache-layer state tensors to tensors before writing the file.
Useful? React with 👍 / 👎.
| DataType::U32 => Ok(Dtype::U32), | ||
| DataType::I64 => Ok(Dtype::I64), | ||
| DataType::U64 => Ok(Dtype::U64), | ||
| DataType::F64 | DataType::I4 | DataType::U4 | DataType::I16 | DataType::U16 => Err(dtype), |
There was a problem hiding this comment.
When callers use the new public write_safetensors helper with SafeTensorData whose data_type is DataType::I16, DataType::U16, or DataType::F64, this arm returns InvalidInput even though the safetensors dtype enum in this module has corresponding I16/U16/F64 variants and ArrayElement supports those element types. That makes otherwise valid tensors unexportable through this writer; only the sub-byte I4/U4 cases need to be rejected here.
Useful? React with 👍 / 👎.
| let request = json!({ | ||
| "add_generation_prompt": true, | ||
| "messages": messages.iter().map(|message| message.resolve(&token_codec_config)).collect::<Vec<_>>(), | ||
| "bos_token": token_codec_config.bos_token, | ||
| "enable_thinking": true, | ||
| }); |
There was a problem hiding this comment.
For chat templates that reference eos_token while rendering this trace, InputProcessorDefault::process receives token_codec_config.eos_token and the exported rendered_request/tokens are produced with that value, but the saved request metadata only carries bos_token. Consumers that replay the saved request with the saved prompt_template will either render with an undefined/null EOS token or fail to reproduce the recorded prompt, so include eos_token alongside the other template inputs.
Useful? React with 👍 / 👎.
ry2009
force-pushed
the
ryan/trace-safetensors-export
branch
from
8771473 to
901ad9f
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 901ad9fa50
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if dim != source_dim || history != source_history { | ||
| return Err(Error::InvalidModelConfig(format!( |
There was a problem hiding this comment.
For Mamba/DeltaNet/ShortConv traces where Lalamo records updated_state.conv_state as [1, suffix_len, dim], this equality check rejects any prompt whose suffix is longer than the native convolution history (kernel_size - 1). The previous SsmConvState validation path handled that schema by taking the last history window, so these common traces now fail export with InvalidModelConfig instead of producing a safetensors file.
Useful? React with 👍 / 👎.
ry2009
force-pushed
the
ryan/trace-safetensors-export
branch
from
901ad9f to
44e6b4f
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 44e6b4f526
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| context, | ||
| trace_metadata, | ||
| format!("activation_trace.layer_results.{index}.updated_state.conv_state"), | ||
| &short_conv.conv_state, |
There was a problem hiding this comment.
For ShortConv layers with a prompt/suffix longer than one token, this exports conv_state immediately after run_llm_trace, but ShortConvMixer::run_trie_conv writes the per-token post-states into suffix_state and only updates conv_state later via commit_from_suffix_state_if_valid; this export path never calls that commit. In that common multi-token trace scenario, the saved updated_state.conv_state is the stale pre-trace state rather than the final state after the prompt, so replay/comparison will see an incorrect cache state.
Useful? React with 👍 / 👎.
ry2009
force-pushed
the
ryan/trace-safetensors-export
branch
from
44e6b4f to
05270e9
Compare
ry2009
force-pushed
the
ryan/trace-safetensors-export
branch
from
05270e9 to
4cd5210
Compare
1 participant
Summary
Reworks the Uzu tracing integration around activation/tokenization export instead of local comparison:
traces.safetensors-compatible safetensors outputactivation_trace.token_ids/activation_trace.token_positionsplus rendered request metadataThis pairs with the Lalamo comparator PR: trymirai/lalamo#285
Validation
cargo test -p backend-uzu test_safetensors_metadata_writer --test integrationcargo check -p backend-uzu --testscargo fmt --check -p backend-uzugit diff --checkcargo fmt --checkpasses but prints existing stable-rustfmt warnings about nightly-only config keys.