[Enhance] add model internal metrics #1271
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
nil0x9
changed the title
HAOCHENYE
reviewed
Nov 17, 2025
| device=hidden_states.device) | ||
|
|
||
| attn_output: torch.Tensor = self.attn_impl_func( # type: ignore | ||
| attn_output, extra_info = self.attn_impl_func( # type: ignore |
Collaborator
There was a problem hiding this comment.
It's not a good idea to change the attention function signature like this. Instead, we should define an AttnOutput type (using TypedDict or namedtuple) to represent the attention result."
| ctx.cu_seqlen = cu_seqlen | ||
|
|
||
| return o | ||
| return o, lse |
Collaborator
There was a problem hiding this comment.
The lse should only be returned when specific flags are passed in. Otherwise, only the attention output should be returned. This makes for a cleaner attention interface.
|
|
||
| def register_attn_extra_info_hook(self, module, layer_name=None): | ||
| def hook(module, input, output): | ||
| extra_info = output[1] |
Collaborator
There was a problem hiding this comment.
Suggest refactoring with the AttnOutput
|
|
||
| # do dummy forward to get metrics | ||
| for i in range(0, len(data_batches), self.intra_layer_micro_batch): | ||
| data_batch = data_batches[i : i + self.intra_layer_micro_batch] |
Collaborator
There was a problem hiding this comment.
Is it necessary to consider intra_layer_micro_batch?
HAOCHENYE
force-pushed
the
dev-add-model-internal-metrics
branch
from
f4014fc to
97f1564
Compare
nil0x9
force-pushed
the
dev-add-model-internal-metrics
branch
9 times, most recently
from
e557bee to
39f7de0
Compare
1. attn logits metrics (max lse / max logits) 2. add moe metrics (mean & max router logits, maxvio & drop ratio) 3. model layerwise weight rms norm 4. est global batch tokens
…ce it is no longer needed
[Enhance] refactor internal metrics to use TypedDict
[Fix] fix rms_norm no_grad
nil0x9
force-pushed
the
dev-add-model-internal-metrics
branch
from
39f7de0 to
0a5f008
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
No description provided.