From 20a15a54848be2bff82d978aa36c83f6a2d6e511 Mon Sep 17 00:00:00 2001 From: Jason Chiu Date: Thu, 20 Feb 2025 03:12:49 +0000 Subject: [PATCH] fix: return float instead of tensor from `get_rotary_seq_len` --- megatron/core/models/common/embeddings/rotary_pos_embedding.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py index e3923d016b..2d633962eb 100644 --- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py +++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py @@ -198,7 +198,7 @@ def get_rotary_seq_len( if packed_seq_params is not None: # max_seqlen are the max sequence length in the packed sequence before being divived # by the tp and cp size. - return max(packed_seq_params.max_seqlen_q, packed_seq_params.max_seqlen_kv) + return max(packed_seq_params.max_seqlen_q, packed_seq_params.max_seqlen_kv).item() elif inference_params is not None: rotary_seq_len = inference_params.max_sequence_length else: