Update data collator to support sequence_length #38536
Conversation
There was a problem hiding this comment.
This looks good, but we might need a test to verify it's working okay and to ensure the behaviour doesn't regress in future updates!
Also, this is a nit but we probably don't need the is_sequence_length_provided bool when it's only used once, you can just test the condition directly
| @@ -2028,7 +2038,14 @@ def __call__(self, features, return_tensors=None, separator_id=None): | |||
| else: | |||
| batch["labels"] += [separator_id] + input_ids[1:] | |||
| if self.return_position_ids: | |||
| batch["position_ids"] += list(range(len(input_ids))) | |||
| if is_sequence_length_provided: | |||
There was a problem hiding this comment.
| if is_sequence_length_provided: | |
| if "sequence_length" in features[0]: |
There was a problem hiding this comment.
same logic applies to is_labels_provided variable -- i find current implementation cleaner?
|
Please keep this PR on hold. Actually we are still unsure if it's actually needed |
|
Looks like this might not be needed, @qgallouedec made it simplier by creating position_ids directly during packing and re-using existing functionality of DataCollatorWithFlattening |
What does this PR do?
Adds ability to pass "sequence_length" key to data collator, achieving true absense of cross-document contamination when using packing.
This approach allows to set contiguous position_ids which in turn gets processed by flash attention kernel inside the Attention, allowing for block-diagonal attention masks and improving both speed and quality.
See huggingface/trl#3526 for loss plots
@ArthurZucker @muellerzr @qgallouedec