Skip to content

Fix Incorrect Hidden State Extraction with Right Padding #26

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xiuyuz wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix Incorrect Hidden State Extraction with Right Padding #26

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 20 additions & 3 deletions models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -310,8 +310,19 @@ def generate_latent_batch(
)
past = outputs.past_key_values

e_t = outputs.hidden_states[0][:, -1, :] # [B, D]
last_hidden = outputs.hidden_states[-1][:, -1, :] # [B, D]
# If past_key_values is None, we might have padding.
if past_key_values is None:
# Identify last token index
# attention_mask (at this point, if past is None, it is just original mask)
last_token_indices = attention_mask.sum(1) - 1
Copy link

Copilot AI Dec 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The attention_mask sum could include padding tokens (zeros) in the mask. If a sequence is completely padded (all zeros in attention_mask), then attention_mask.sum(1) would be 0, resulting in last_token_indices being -1. While this is technically a valid Python index (referring to the last element), it would give incorrect behavior for completely padded sequences. Consider adding validation or handling for this edge case.

Suggested change
last_token_indices = attention_mask.sum(1) - 1
token_counts = attention_mask.sum(1)
# Clamp to avoid negative indices when a sequence is fully padded (sum == 0).
last_token_indices = torch.clamp(token_counts - 1, min=0)

Copilot uses AI. Check for mistakes.
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
Copy link

Copilot AI Dec 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For clarity and consistency with how attention_mask is handled, batch_indices should be created using self.device instead of input_ids.device. While input_ids.device should be the same as self.device (otherwise the model call would fail), using self.device explicitly makes the device management more clear and matches the pattern used for attention_mask at line 289.

Suggested change
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
batch_indices = torch.arange(input_ids.shape[0], device=self.device)

Copilot uses AI. Check for mistakes.
e_t = outputs.hidden_states[0][batch_indices, last_token_indices, :]
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]
else:
# Assume no padding in incremental decoding steps
e_t = outputs.hidden_states[0][:, -1, :]
last_hidden = outputs.hidden_states[-1][:, -1, :]
Comment on lines +315 to +324
Copy link

Copilot AI Dec 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The indentation inside this if block appears to use 5 spaces instead of the standard 4 spaces used throughout the rest of the codebase. Lines 315-320 should be indented with 12 spaces (8 base + 4 for the if block) rather than 13 spaces.

Suggested change
# Identify last token index
# attention_mask (at this point, if past is None, it is just original mask)
last_token_indices = attention_mask.sum(1) - 1
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
e_t = outputs.hidden_states[0][batch_indices, last_token_indices, :]
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]
else:
# Assume no padding in incremental decoding steps
e_t = outputs.hidden_states[0][:, -1, :]
last_hidden = outputs.hidden_states[-1][:, -1, :]
# Identify last token index
# attention_mask (at this point, if past is None, it is just original mask)
last_token_indices = attention_mask.sum(1) - 1
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
e_t = outputs.hidden_states[0][batch_indices, last_token_indices, :]
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]
else:
# Assume no padding in incremental decoding steps
e_t = outputs.hidden_states[0][:, -1, :]
last_hidden = outputs.hidden_states[-1][:, -1, :]

Copilot uses AI. Check for mistakes.
Comment on lines +315 to +324
Copy link

Copilot AI Dec 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The indentation inside this if block appears to use 5 spaces instead of the standard 4 spaces used throughout the rest of the codebase. Lines 323-324 should be indented with 12 spaces (8 base + 4 for the else block) rather than 13 spaces.

Suggested change
# Identify last token index
# attention_mask (at this point, if past is None, it is just original mask)
last_token_indices = attention_mask.sum(1) - 1
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
e_t = outputs.hidden_states[0][batch_indices, last_token_indices, :]
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]
else:
# Assume no padding in incremental decoding steps
e_t = outputs.hidden_states[0][:, -1, :]
last_hidden = outputs.hidden_states[-1][:, -1, :]
# Identify last token index
# attention_mask (at this point, if past is None, it is just original mask)
last_token_indices = attention_mask.sum(1) - 1
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
e_t = outputs.hidden_states[0][batch_indices, last_token_indices, :]
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]
else:
# Assume no padding in incremental decoding steps
e_t = outputs.hidden_states[0][:, -1, :]
last_hidden = outputs.hidden_states[-1][:, -1, :]

Copilot uses AI. Check for mistakes.

h_t = last_hidden.detach().clone()

e_t_plus_1 = None
Expand Down Expand Up @@ -382,7 +393,13 @@ def generate_latent_batch_hidden_state(
return_dict=True,
)
past = outputs.past_key_values
last_hidden = outputs.hidden_states[-1][:, -1, :]

if past_key_values is None:
last_token_indices = attention_mask.sum(1) - 1
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
Copy link

Copilot AI Dec 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For clarity and consistency with how attention_mask is handled, batch_indices should be created using self.HF_device instead of input_ids.device. While input_ids.device should be the same as self.HF_device (otherwise the model call would fail), using self.HF_device explicitly makes the device management more clear and matches the pattern used for attention_mask at line 375.

Suggested change
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
batch_indices = torch.arange(input_ids.shape[0], device=self.HF_device)

Copilot uses AI. Check for mistakes.
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]
Comment on lines +398 to +400
Copy link

Copilot AI Dec 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The attention_mask sum could include padding tokens (zeros) in the mask. If a sequence is completely padded (all zeros in attention_mask), then attention_mask.sum(1) would be 0, resulting in last_token_indices being -1. While this is technically a valid Python index (referring to the last element), it would give incorrect behavior for completely padded sequences. Consider adding validation or handling for this edge case.

Suggested change
last_token_indices = attention_mask.sum(1) - 1
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]
seq_lengths = attention_mask.sum(1)
# Handle fully padded sequences (sum == 0) to avoid negative indices (-1)
if torch.any(seq_lengths == 0):
seq_lengths = seq_lengths.clone()
seq_lengths[seq_lengths == 0] = 1
last_token_indices = seq_lengths - 1
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]

Copilot uses AI. Check for mistakes.
else:
last_hidden = outputs.hidden_states[-1][:, -1, :]
Comment on lines +398 to +402
Copy link

Copilot AI Dec 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The indentation inside this else block appears to use 5 spaces instead of the standard 4 spaces used throughout the rest of the codebase. Line 402 should be indented with 12 spaces (8 base + 4 for the else block) rather than 13 spaces.

Suggested change
last_token_indices = attention_mask.sum(1) - 1
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]
else:
last_hidden = outputs.hidden_states[-1][:, -1, :]
last_token_indices = attention_mask.sum(1) - 1
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]
else:
last_hidden = outputs.hidden_states[-1][:, -1, :]

Copilot uses AI. Check for mistakes.
Comment on lines +398 to +402
Copy link

Copilot AI Dec 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The indentation inside this if block appears to use 5 spaces instead of the standard 4 spaces used throughout the rest of the codebase. Lines 398-400 should be indented with 12 spaces (8 base + 4 for the if block) rather than 13 spaces.

Suggested change
last_token_indices = attention_mask.sum(1) - 1
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]
else:
last_hidden = outputs.hidden_states[-1][:, -1, :]
last_token_indices = attention_mask.sum(1) - 1
batch_indices = torch.arange(input_ids.shape[0], device=input_ids.device)
last_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices, :]
else:
last_hidden = outputs.hidden_states[-1][:, -1, :]

Copilot uses AI. Check for mistakes.

curr_output_embedding = []
curr_output_embedding.append(outputs.hidden_states[0]) # input embedding
Expand Down