(
weights: torch.Tensor, # [V, D] (may be a TP shard over dim V)
hidden_states: torch.Tensor, # [n_hidden_states, D]
num_samples: int,
temperature: torch.Tensor, # scalar (0-d)
return_probs: bool = False,
seed: int = None,
tl_matmul: bool = False,
top_k: int | None = None,
top_p: float | None = None,
use_qitra: bool = False,
tp: "TPInfo" = TP1,
)
| 27 | |
| 28 | |
| 29 | def sample( |
| 30 | weights: torch.Tensor, # [V, D] (may be a TP shard over dim V) |
| 31 | hidden_states: torch.Tensor, # [n_hidden_states, D] |
| 32 | num_samples: int, |
| 33 | temperature: torch.Tensor, # scalar (0-d) |
| 34 | return_probs: bool = False, |
| 35 | seed: int = None, |
| 36 | tl_matmul: bool = False, |
| 37 | top_k: int | None = None, |
| 38 | top_p: float | None = None, |
| 39 | use_qitra: bool = False, |
| 40 | tp: "TPInfo" = TP1, |
| 41 | ): |
| 42 | if seed is not None: |
| 43 | torch.manual_seed(seed) |
| 44 | if tl_matmul: |
| 45 | logits = matmul(hidden_states, weights) # [n_hidden_states, V] |
| 46 | else: |
| 47 | logits = hidden_states @ weights.T # [n_hidden_states, V] |
| 48 | if tp.size > 1: |
| 49 | logits = _allgather_logits(logits) # shape [H, V_local] -> [H, V] |
| 50 | # Upcast to float32 before temperature scaling: Qitra asserts float32, and |
| 51 | # torch.multinomial produces imprecise distributions with bfloat16. |
| 52 | # See findings/upcasting-before-softmax.md. |
| 53 | logits = logits.float() / temperature |
| 54 | if use_qitra: |
| 55 | probs = apply_top_k_top_p_qitra(logits, top_k, top_p) |
| 56 | else: |
| 57 | probs = apply_top_k_top_p(logits, top_k, top_p) |
| 58 | samples = torch.multinomial(probs, num_samples, replacement=True) |
| 59 | if return_probs: |
| 60 | return samples, probs |
| 61 | return samples |
| 62 | |
| 63 | |
| 64 | def _fast_multinomial(probs: torch.Tensor, num_samples: int) -> torch.Tensor: |
no test coverage detected