Skip to content

make the PointerHead ready for half precision [DONT MERGE] #154

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
ArneBinder wants to merge 10 commits into
base: main
Choose a base branch
from
Open

make the PointerHead ready for half precision [DONT MERGE] #154

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
7ecdc60
use min value of dtype, adjust test and add new test case
ArneBinder Feb 6, 2025
eeda275
use different min_float_vals
ArneBinder Feb 6, 2025
79fec01
set dtype when new_full
ArneBinder Feb 6, 2025
1b06f96
ad prints
ArneBinder Feb 6, 2025
5291d98
remove prints
ArneBinder Feb 6, 2025
8e21ef4
cleanup
ArneBinder Feb 6, 2025
2e4fd26
adjust tests for BartAsPointerNetwork
ArneBinder Feb 6, 2025
520d6ab
temporarily hardcode torch_dtype=float16 for SimpleGenerativeModel
ArneBinder Feb 6, 2025
acf722a
revert: temporarily hardcode torch_dtype=float16 for SimpleGenerative…
ArneBinder Feb 6, 2025
de975b5
adjust tests for BartAsPointerNetwork
ArneBinder Feb 6, 2025
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
14 changes: 10 additions & 4 deletions src/pie_modules/models/components/pointer_head.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -261,14 +261,16 @@ def forward(
decoder_attention_mask: Optional[torch.LongTensor] = None,
constraints: Optional[torch.LongTensor] = None,
):
min_float_val = torch.finfo(last_hidden_state.dtype).min
# assemble the logits
logits = last_hidden_state.new_full(
(
last_hidden_state.size(0),
last_hidden_state.size(1),
self.pointer_offset + encoder_input_ids.size(-1),
),
fill_value=-1e24,
fill_value=min_float_val,
dtype=last_hidden_state.dtype,
)

# eos and label scores depend only on the decoder output
Expand All @@ -295,7 +297,8 @@ def forward(
# never point to the padding or the eos token in the encoder input
# TODO: why not excluding the bos token? seems to give worse results, but not tested extensively
mask_invalid = encoder_attention_mask.eq(0) | encoder_input_ids.eq(self.eos_token_id)
avg_word_scores = avg_word_scores.masked_fill(mask_invalid.unsqueeze(1), -1e32)
min_float_val = torch.finfo(avg_word_scores.dtype).min
avg_word_scores = avg_word_scores.masked_fill(mask_invalid.unsqueeze(1), min_float_val)

# Note: the remaining row in logits contains the score for the bos token which should be never generated!
logits[:, :, [self.eos_id]] = eos_scores
Expand Down Expand Up @@ -331,13 +334,15 @@ def forward(
constraints_word_scores = torch.einsum(
"blh,bnh->bln", last_hidden_state, constraints_src_outputs
)
min_float_val = torch.finfo(last_hidden_state.dtype).min
constraints_logits = last_hidden_state.new_full(
(
last_hidden_state.size(0),
last_hidden_state.size(1),
self.pointer_offset + encoder_input_ids.size(-1),
),
fill_value=-1e24,
fill_value=min_float_val,
dtype=last_hidden_state.dtype,
)
constraints_logits[:, :, self.label_ids] = constraints_label_scores
constraints_logits[:, :, self.pointer_offset :] = constraints_word_scores
Expand All @@ -346,7 +351,8 @@ def forward(
constraints_logits_valid = constraints_logits[mask]
constraints_valid = constraints[mask]
loss_c = F.binary_cross_entropy(
torch.sigmoid(constraints_logits_valid), constraints_valid.float()
torch.sigmoid(constraints_logits_valid),
constraints_valid.to(dtype=constraints_logits_valid.dtype),
)

if loss is None:
Expand Down
Loading