File size: 1,497 Bytes
84a4b93 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | """Pack OLMo's document offsets for FLA; no global forward patch."""
import torch
def flatten_documents(x: torch.Tensor, cu_doc_lens: torch.Tensor | None) -> torch.Tensor:
"""Flatten before BOTH short convolutions and the recurrent kernel.
OLMo's offsets already refer to the row-major flattened batch. Returning a
view preserves autograd and lets the caller restore the original shape.
"""
if x.ndim != 3:
raise ValueError("Expected x with shape [batch, sequence, hidden]")
if cu_doc_lens is None:
return x
if cu_doc_lens.ndim != 1 or cu_doc_lens.numel() < 2:
raise ValueError("cu_doc_lens must be a one-dimensional offsets tensor")
if cu_doc_lens.dtype not in (torch.int32, torch.int64):
raise ValueError("cu_doc_lens must have integer dtype int32 or int64")
if cu_doc_lens.device != x.device:
raise ValueError("cu_doc_lens and x must be on the same device")
batch, length, hidden = x.shape
# Avoid per-layer host synchronization and graph breaks under torch.compile.
if not torch.compiler.is_compiling():
if int(cu_doc_lens[0]) != 0 or int(cu_doc_lens[-1]) != batch * length:
raise ValueError("cu_doc_lens must start at 0 and end at batch * sequence")
if not bool((cu_doc_lens[1:] > cu_doc_lens[:-1]).all()):
raise ValueError("cu_doc_lens must increase strictly (no empty documents)")
return x.reshape(1, batch * length, hidden)
|