[Auto Sync] Update logits_processor.py (20260209) (#18503)
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
github-actions[bot]
parent
01e3f4682e
commit
ce95f203b0
@@ -247,6 +247,7 @@ class LogitsProcessor(nn.Module):
|
|||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.config = config
|
self.config = config
|
||||||
|
self.vocab_size = config.vocab_size
|
||||||
self.logit_scale = logit_scale
|
self.logit_scale = logit_scale
|
||||||
self.use_attn_tp_group = get_global_server_args().enable_dp_lm_head
|
self.use_attn_tp_group = get_global_server_args().enable_dp_lm_head
|
||||||
self.use_fp32_lm_head = get_global_server_args().enable_fp32_lm_head
|
self.use_fp32_lm_head = get_global_server_args().enable_fp32_lm_head
|
||||||
@@ -906,23 +907,23 @@ class LogitsProcessor(nn.Module):
|
|||||||
return hidden_states, hidden_states
|
return hidden_states, hidden_states
|
||||||
|
|
||||||
def _gather_attn_tp_logits(self, logits: torch.Tensor) -> torch.Tensor:
|
def _gather_attn_tp_logits(self, logits: torch.Tensor) -> torch.Tensor:
|
||||||
if self.config.vocab_size % self.attn_tp_size == 0:
|
if self.vocab_size % self.attn_tp_size == 0:
|
||||||
global_logits = torch.empty(
|
global_logits = torch.empty(
|
||||||
(
|
(
|
||||||
self.attn_tp_size,
|
self.attn_tp_size,
|
||||||
logits.shape[0],
|
logits.shape[0],
|
||||||
self.config.vocab_size // self.attn_tp_size,
|
self.vocab_size // self.attn_tp_size,
|
||||||
),
|
),
|
||||||
device=logits.device,
|
device=logits.device,
|
||||||
dtype=logits.dtype,
|
dtype=logits.dtype,
|
||||||
)
|
)
|
||||||
attn_tp_all_gather_into_tensor(global_logits, logits)
|
attn_tp_all_gather_into_tensor(global_logits, logits)
|
||||||
global_logits = global_logits.permute(1, 0, 2).reshape(
|
global_logits = global_logits.permute(1, 0, 2).reshape(
|
||||||
logits.shape[0], self.config.vocab_size
|
logits.shape[0], self.vocab_size
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
global_logits = torch.empty(
|
global_logits = torch.empty(
|
||||||
(self.config.vocab_size, logits.shape[0]),
|
(self.vocab_size, logits.shape[0]),
|
||||||
device=logits.device,
|
device=logits.device,
|
||||||
dtype=logits.dtype,
|
dtype=logits.dtype,
|
||||||
)
|
)
|
||||||
@@ -955,10 +956,10 @@ class LogitsProcessor(nn.Module):
|
|||||||
if logits_metadata.next_token_logits_buffer is not None:
|
if logits_metadata.next_token_logits_buffer is not None:
|
||||||
logits_buffer = logits_metadata.next_token_logits_buffer
|
logits_buffer = logits_metadata.next_token_logits_buffer
|
||||||
assert logits_buffer.dtype == torch.float
|
assert logits_buffer.dtype == torch.float
|
||||||
logits_buffer.copy_(logits[:, : self.config.vocab_size])
|
logits_buffer.copy_(logits[:, : self.vocab_size])
|
||||||
logits = logits_buffer
|
logits = logits_buffer
|
||||||
else:
|
else:
|
||||||
logits = logits[:, : self.config.vocab_size].float()
|
logits = logits[:, : self.vocab_size].float()
|
||||||
return logits
|
return logits
|
||||||
|
|
||||||
def _get_dllm_logits(
|
def _get_dllm_logits(
|
||||||
|
|||||||
Reference in New Issue
Block a user