diff --git a/rl_engine/kernels/ops/triton/loss/linear_logp.py b/rl_engine/kernels/ops/triton/loss/linear_logp.py index 561d81e4..fca26322 100644 --- a/rl_engine/kernels/ops/triton/loss/linear_logp.py +++ b/rl_engine/kernels/ops/triton/loss/linear_logp.py @@ -9,6 +9,10 @@ import triton.language as tl from rl_engine.kernels.ops.pytorch.loss.linear_logp import ( + _merge_tp_local_logp, + _require_distributed_initialized, + _validate_global_targets, + _validate_tp_vocab_partition_cached, chunked_linear_logp_backward, should_use_tensor_parallel_linear_logp, tensor_parallel_linear_logp, @@ -236,3 +240,642 @@ def apply( "Mask or filter padding / ignore-index values (e.g. -100) before this op." ) return _LinearLogpFunction.apply(hidden, lm_head_weight, bias, target_ids) + + +# --------------------------------------------------------------------------- +# Deterministic (strict-contract) Triton linear_logp -- the portable analogue +# of ``sm90_deterministic_linear_logp`` (see rl_engine/kernels/ops/cuda/loss/ +# linear_logp.py and csrc/cuda/fused_linear_logp_sm90.cu). One source for CUDA +# and ROCm. +# +# Frozen numerical contract (bit-affecting; bump the version to change any): +# * vocab splitting ....... TRITON_N_SPLIT_CONTRACT fixed splits over +# ceil(V / BLOCK_V) tiles, boundaries depend only +# on V -- never on batch shape, occupancy, or CU +# count +# * within a split ........ ascending-v0 online-softmax chain over +# BLOCK_V-wide tiles; fixed BLOCK_D K-chain inside +# tl.dot with IEEE FP32 accumulation +# * cross-split merge ..... ascending-split sequential scalar chains +# * padding lanes ......... -inf to max, exp() -> exact 0 to sum +# * temperature ........... multiplies stats and the selected logit by +# 1/temperature in the FP32 epilogue, never the +# stored logits +# * final clamp ........... logp = min(zt - lse, 0) +# Row-block size (_DET_BLOCK_N) only moves work along the token axis: rows are +# numerically independent, so it is bit-neutral and may be tuned freely. +# --------------------------------------------------------------------------- + +TRITON_LINEAR_LOGP_CONTRACT_VERSION = "triton-fused-linear-logp-contract-v1" +TRITON_N_SPLIT_CONTRACT = 64 + +_DET_BLOCK_N = 32 # bit-neutral row tile +_DET_BLOCK_V = 64 # contract: vocab tile width +_DET_BLOCK_D = 64 # contract: K-chain step + + +@triton.jit +def _det_linear_logp_partial_kernel( + h_ptr, # hidden [N, D] + w_ptr, # lm_head_weight [V, D] + b_ptr, # bias [V] (dummy when HAS_BIAS=False) + t_ptr, # temperature [N] fp32 (dummy when HAS_TEMP=False) + tgt_ptr, # target_ids [N] int64 (global ids) + logits_ptr, # unscaled fp32 logits [N, V] (dummy when STORE_LOGITS=False) + part_max_ptr, # [n_split, N] fp32 + part_sum_ptr, # [n_split, N] fp32 + part_zt_ptr, # [n_split, N] fp32 + N, + D, + V, + n_split, + vocab_start, + real_vocab_end, # local column bound of the real vocabulary + stride_hn, + stride_hd, + stride_wv, + stride_wd, + HAS_BIAS: tl.constexpr, + HAS_TEMP: tl.constexpr, + STORE_LOGITS: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_V: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """One program per (row block, vocab split): fp32 (max, sumexp, zt) partials. + + The split owns a contiguous range of BLOCK_V tiles fixed by (V, n_split) + alone and folds them in ascending order into an online-softmax state, so + the reduction tree for any row never depends on the batch shape.""" + pid = tl.program_id(0) + split = tl.program_id(1) + rows = pid * BLOCK_N + tl.arange(0, BLOCK_N) + row_mask = rows < N + target = tl.load(tgt_ptr + rows, mask=row_mask, other=-1) + + total_vtiles = tl.cdiv(V, BLOCK_V) + vtiles_per_split = tl.cdiv(total_vtiles, n_split) + vt_begin = split * vtiles_per_split + vt_end = tl.minimum(vt_begin + vtiles_per_split, total_vtiles) + + inv_t = tl.full((BLOCK_N,), 1.0, tl.float32) + if HAS_TEMP: + inv_t = 1.0 / tl.load(t_ptr + rows, mask=row_mask, other=1.0) + + m = tl.full((BLOCK_N,), float("-inf"), tl.float32) + s = tl.zeros((BLOCK_N,), tl.float32) + zt = tl.zeros((BLOCK_N,), tl.float32) + + for vt in range(vt_begin, vt_end): + vcols = vt * BLOCK_V + tl.arange(0, BLOCK_V) + vmask = vcols < V + + acc = tl.zeros((BLOCK_N, BLOCK_V), tl.float32) + for d0 in range(0, D, BLOCK_D): + offs_d = d0 + tl.arange(0, BLOCK_D) + d_mask = offs_d < D + h = tl.load( + h_ptr + rows[:, None] * stride_hn + offs_d[None, :] * stride_hd, + mask=row_mask[:, None] & d_mask[None, :], + other=0.0, + ) + w = tl.load( + w_ptr + vcols[:, None] * stride_wv + offs_d[None, :] * stride_wd, + mask=vmask[:, None] & d_mask[None, :], + other=0.0, + ) + acc += tl.dot(h, tl.trans(w), input_precision="ieee") + + if HAS_BIAS: + acc += tl.load(b_ptr + vcols, mask=vmask, other=0.0).to(tl.float32)[None, :] + + real_mask = vmask & (vcols < real_vocab_end) + if STORE_LOGITS: + # [contract] stored logits are unscaled: bias applied, padding -inf, + # temperature never applied. + tl.store( + logits_ptr + rows[:, None].to(tl.int64) * V + vcols[None, :], + tl.where(real_mask[None, :], acc, float("-inf")), + mask=row_mask[:, None] & vmask[None, :], + ) + + val = tl.where(real_mask[None, :], acc, float("-inf")) + if HAS_TEMP: + val = val * inv_t[:, None] + + is_target = (vcols[None, :].to(tl.int64) + vocab_start) == target[:, None] + zt += tl.sum(tl.where(is_target & real_mask[None, :], val, 0.0), axis=1) + + tile_max = tl.max(val, axis=1) + new_m = tl.maximum(m, tile_max) + finite = new_m != float("-inf") + alpha = tl.where(finite, tl.exp(m - new_m), 1.0) + p = tl.where(finite[:, None], tl.exp(val - new_m[:, None]), 0.0) + s = s * alpha + tl.sum(p, axis=1) + m = new_m + + base = split.to(tl.int64) * N + rows + tl.store(part_max_ptr + base, m, mask=row_mask) + tl.store(part_sum_ptr + base, s, mask=row_mask) + tl.store(part_zt_ptr + base, zt, mask=row_mask) + + +@triton.jit +def _det_linear_logp_merge_kernel( + part_max_ptr, # [n_split, N] fp32 + part_sum_ptr, + part_zt_ptr, + zt_ptr, # output [N] fp32: selected (scaled) target logit + lse_ptr, # output [N] fp32 + N, + n_split, + BLOCK_N: tl.constexpr, +): + """[contract] Ascending-split sequential merge of the fp32 partials.""" + pid = tl.program_id(0) + rows = pid * BLOCK_N + tl.arange(0, BLOCK_N) + row_mask = rows < N + + rows64 = rows.to(tl.int64) + m = tl.load(part_max_ptr + rows, mask=row_mask, other=float("-inf")) + for split in range(1, n_split): + base = rows64 + split * N + m = tl.maximum(m, tl.load(part_max_ptr + base, mask=row_mask, other=float("-inf"))) + + finite = m != float("-inf") + s = tl.zeros((BLOCK_N,), tl.float32) + zt = tl.zeros((BLOCK_N,), tl.float32) + for split in range(0, n_split): + base = rows64 + split * N + pm = tl.load(part_max_ptr + base, mask=row_mask, other=float("-inf")) + ps = tl.load(part_sum_ptr + base, mask=row_mask, other=0.0) + term = tl.where(finite & (pm != float("-inf")), ps * tl.exp(pm - m), 0.0) + s = s + term + zt = zt + tl.load(part_zt_ptr + base, mask=row_mask, other=0.0) + + lse = m + tl.log(s) + tl.store(zt_ptr + rows, zt, mask=row_mask) + tl.store(lse_ptr + rows, lse, mask=row_mask) + + +def _det_linear_logp_local( + hidden_2d: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + target_1d: torch.Tensor, + temperature: Optional[torch.Tensor], + *, + vocab_start: int, + real_vocab_end: int, + store_logits: bool, +) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Run the frozen-contract kernels on one weight shard. + + Returns ``(zt, lse, logits)``: the temperature-scaled selected-target + contribution (0 when this shard does not own the target), the shard-local + log-sum-exp over its real-vocabulary columns (``-inf`` for an all-padding + shard), and the unscaled fp32 logits when requested.""" + n, d = hidden_2d.shape + v = weight.shape[0] + device = hidden_2d.device + n_split = TRITON_N_SPLIT_CONTRACT + part = torch.empty(3, n_split, max(n, 1), device=device, dtype=torch.float32) + zt = torch.empty(n, device=device, dtype=torch.float32) + lse = torch.empty(n, device=device, dtype=torch.float32) + logits = torch.empty(n, v, device=device, dtype=torch.float32) if store_logits else None + if n == 0: + return zt, lse, logits + dummy = hidden_2d + grid = (triton.cdiv(n, _DET_BLOCK_N), n_split) + _det_linear_logp_partial_kernel[grid]( + hidden_2d, + weight, + bias if bias is not None else dummy, + temperature if temperature is not None else dummy, + target_1d, + logits if logits is not None else dummy, + part[0], + part[1], + part[2], + n, + d, + v, + n_split, + int(vocab_start), + int(real_vocab_end), + hidden_2d.stride(0), + hidden_2d.stride(1), + weight.stride(0), + weight.stride(1), + HAS_BIAS=bias is not None, + HAS_TEMP=temperature is not None, + STORE_LOGITS=store_logits, + BLOCK_N=_DET_BLOCK_N, + BLOCK_V=_DET_BLOCK_V, + BLOCK_D=_DET_BLOCK_D, + ) + _det_linear_logp_merge_kernel[(triton.cdiv(n, 256),)]( + part[0], + part[1], + part[2], + zt, + lse, + n, + n_split, + BLOCK_N=256, + ) + return zt, lse, logits + + +def _det_prepare( + hidden: torch.Tensor, + weight: torch.Tensor, + target_ids: torch.Tensor, + temperature: Optional[torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + if hidden.device.type not in ("cuda", "hip", "xpu"): + raise RuntimeError( + "deterministic Triton linear_logp requires a GPU tensor, got " + f"device '{hidden.device}'." + ) + if hidden.dim() < 1: + raise ValueError("hidden must have at least one dimension") + if weight.size(-1) != hidden.size(-1): + raise ValueError( + f"hidden dim {hidden.size(-1)} must match lm_head_weight dim {weight.size(-1)}" + ) + hidden_2d = hidden.reshape(-1, hidden.size(-1)).contiguous() + weight_c = weight.contiguous() + target_1d = target_ids.reshape(-1).to(device=hidden_2d.device, dtype=torch.long).contiguous() + if target_1d.numel() != hidden_2d.size(0): + raise ValueError("target_ids must have one id per hidden row") + temp_arg = None + if temperature is not None: + temp_arg = temperature.to(device=hidden_2d.device, dtype=torch.float32).reshape(-1) + if temp_arg.numel() == 1: + temp_arg = temp_arg.expand(hidden_2d.size(0)).contiguous() + else: + temp_arg = temp_arg.contiguous() + if temp_arg.numel() != hidden_2d.size(0) or bool((temp_arg <= 0).any().item()): + raise ValueError("temperature must be positive and scalar or per-token") + return hidden_2d, weight_c, target_1d, temp_arg + + +def _det_reference_dlogits( + hidden_2d: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + target_local: torch.Tensor, + owned: Optional[torch.Tensor], + lse: torch.Tensor, + temp: Optional[torch.Tensor], + grad_logp: Optional[torch.Tensor], + grad_lse: Optional[torch.Tensor], + *, + real_vocab_end: int, +) -> torch.Tensor: + """FP32 recompute of d(logits) for the strict backward (local shard). + + ``d(logp)/d(logits) = onehot(target) - softmax`` and + ``d(lse)/d(logits) = softmax`` on the temperature-scaled logits; the chain + through the scaling divides by the temperature once more at the end.""" + hidden_f = hidden_2d.float() + weight_f = weight.float() + logits = torch.nn.functional.linear( + hidden_f, weight_f, bias.float() if bias is not None else None + ) + if temp is not None: + logits = logits / temp.reshape(-1, 1) + if real_vocab_end < logits.size(1): + columns = torch.arange(logits.size(1), device=logits.device) + logits = logits.masked_fill(columns[None, :] >= real_vocab_end, float("-inf")) + lse_f = lse.reshape(-1, 1).float() + probs = torch.exp(logits - lse_f) + logp_grad = torch.zeros_like(lse) if grad_logp is None else grad_logp.reshape(-1).float() + lse_grad = torch.zeros_like(lse) if grad_lse is None else grad_lse.reshape(-1).float() + dlogits = probs * (lse_grad - logp_grad).reshape(-1, 1) + rows = torch.arange(target_local.numel(), device=target_local.device) + if owned is None: + dlogits[rows, target_local] += logp_grad + else: + hit = rows[owned] + dlogits[hit, target_local[owned]] += logp_grad[owned] + if temp is not None: + dlogits = dlogits / temp.reshape(-1, 1) + return dlogits + + +class _DetTritonLinearLogpAutograd(torch.autograd.Function): + @staticmethod + def forward(ctx, hidden, weight, bias, target, temperature, real_vocab_size): + hidden_2d, weight_c, target_1d, temp = _det_prepare(hidden, weight, target, temperature) + real_vocab = weight_c.size(0) if int(real_vocab_size) < 0 else int(real_vocab_size) + zt, lse, _ = _det_linear_logp_local( + hidden_2d, + weight_c, + bias.contiguous() if bias is not None else None, + target_1d, + temp, + vocab_start=0, + real_vocab_end=real_vocab, + store_logits=False, + ) + logp = torch.minimum(zt - lse, torch.zeros_like(lse)) + ctx.save_for_backward( + hidden_2d, + weight_c, + bias.contiguous() if bias is not None else hidden_2d.new_empty(0), + target_1d, + lse, + temp if temp is not None else hidden_2d.new_empty(0), + ) + ctx.has_bias = bias is not None + ctx.real_vocab_end = real_vocab + ctx.lead_shape = hidden.shape[:-1] + ctx.hidden_dtype = hidden.dtype + ctx.weight_dtype = weight.dtype + ctx.bias_dtype = bias.dtype if bias is not None else None + return logp.reshape(ctx.lead_shape), lse.reshape(ctx.lead_shape) + + @staticmethod + def backward(ctx, grad_logp, grad_lse): + if grad_logp is None and grad_lse is None: + return (None, None, None, None, None, None) + hidden_2d, weight, bias, target, lse, temp = ctx.saved_tensors + with torch.no_grad(): + dlogits = _det_reference_dlogits( + hidden_2d, + weight, + bias if ctx.has_bias else None, + target, + None, + lse, + temp if temp.numel() else None, + grad_logp, + grad_lse, + real_vocab_end=ctx.real_vocab_end, + ) + grad_hidden = dlogits.matmul(weight.float()) + grad_weight = dlogits.transpose(0, 1).matmul(hidden_2d.float()) + grad_bias = dlogits.sum(0) if ctx.has_bias else None + return ( + grad_hidden.reshape((*tuple(ctx.lead_shape), weight.size(1))).to(ctx.hidden_dtype), + grad_weight.to(ctx.weight_dtype), + None if grad_bias is None else grad_bias.to(ctx.bias_dtype), + None, + None, + None, + ) + + +def triton_deterministic_linear_logp( + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + temperature: Optional[torch.Tensor] = None, + return_logits: bool = False, + real_vocab_size: int = -1, +): + """Direct forward entry into the frozen-contract Triton fused linear_logp. + + Portable strict-mode boundary (CUDA and ROCm from one kernel source): two + callers with byte-identical ``hidden`` / ``lm_head_weight`` / + ``target_ids`` / ``temperature`` and the same + ``TRITON_LINEAR_LOGP_CONTRACT_VERSION`` receive bitwise-identical FP32 + ``(logp, lse)`` -- for any batch shape that contains the row. + + * ``temperature`` (float32, scalar or one value per token, strictly + positive) scales the streaming-softmax stats and the selected logit in + the FP32 epilogue -- the same rounding point on every caller. + * ``return_logits=True`` additionally returns the unscaled FP32 logits + (bias applied, padded lanes ``-inf``) without changing the + ``logp``/``lse`` bytes. + * ``real_vocab_size`` masks padded vocabulary lanes out of the LSE: lanes + at or beyond it contribute exactly ``-inf``/``0``. ``-1`` disables + masking. + """ + hidden_2d, weight, target_1d, temp_arg = _det_prepare( + hidden, lm_head_weight, target_ids, temperature + ) + lead_shape = hidden.shape[:-1] + real_vocab = weight.size(0) if int(real_vocab_size) < 0 else int(real_vocab_size) + if not 0 < real_vocab <= weight.size(0): + raise ValueError(f"real_vocab_size must be in [1, {weight.size(0)}], got {real_vocab_size}") + if target_1d.numel() and bool(((target_1d < 0) | (target_1d >= real_vocab)).any().item()): + target_min = int(target_1d.min().item()) + target_max = int(target_1d.max().item()) + raise ValueError( + f"target_ids must be in the real vocabulary [0, {real_vocab}), " + f"got [{target_min}, {target_max}]" + ) + if ( + torch.is_grad_enabled() + and ( + hidden.requires_grad + or lm_head_weight.requires_grad + or (bias is not None and bias.requires_grad) + ) + and not return_logits + ): + return _DetTritonLinearLogpAutograd.apply( + hidden, lm_head_weight, bias, target_ids, temp_arg, int(real_vocab_size) + ) + zt, lse, logits = _det_linear_logp_local( + hidden_2d, + weight, + bias.contiguous() if bias is not None else None, + target_1d, + temp_arg, + vocab_start=0, + real_vocab_end=real_vocab, + store_logits=return_logits, + ) + logp = torch.minimum(zt - lse, torch.zeros_like(lse)).reshape(lead_shape) + lse = lse.reshape(lead_shape) + if return_logits: + assert logits is not None + return logp, lse, logits.reshape(*lead_shape, weight.size(0)) + return logp, lse + + +def _det_tp_run( + hidden, + weight, + bias, + target, + *, + vocab_start, + global_vocab, + real_vocab, + temperature, + tp_group, +): + hidden_2d, weight_c, target_1d, temp = _det_prepare(hidden, weight, target, temperature) + global_vocab = _validate_tp_vocab_partition_cached( + tp_group=tp_group, + device=hidden_2d.device, + vocab_start_index=int(vocab_start), + local_vocab_size=weight_c.size(0), + global_vocab_size=int(global_vocab), + ) + if not 0 < int(real_vocab) <= global_vocab: + raise ValueError(f"invalid real_vocab_size={real_vocab} for padded vocab={global_vocab}") + _validate_global_targets(target_1d, int(real_vocab), tp_group) + dist = _require_distributed_initialized() + owners = ( + (target_1d >= int(vocab_start)) & (target_1d < int(vocab_start) + weight_c.size(0)) + ).to(torch.int32) + dist.all_reduce(owners, op=dist.ReduceOp.SUM, group=tp_group) + if bool((owners != 1).any().item()): + raise ValueError("each selected target must have exactly one TP LM-head owner") + local_real_end = max(0, min(weight_c.size(0), int(real_vocab) - int(vocab_start))) + local_zt, local_lse, _ = _det_linear_logp_local( + hidden_2d, + weight_c, + bias.contiguous() if bias is not None else None, + target_1d, + temp, + vocab_start=int(vocab_start), + real_vocab_end=local_real_end, + store_logits=False, + ) + # [contract] Rank merge is the shared explicit ascending-rank chain. + logp, lse = _merge_tp_local_logp(local_lse, local_zt, tp_group=tp_group) + return logp, lse, hidden_2d, weight_c, target_1d, temp + + +class _DetTritonTensorParallelLinearLogpAutograd(torch.autograd.Function): + @staticmethod + def forward( + ctx, + hidden, + weight, + bias, + target, + vocab_start, + global_vocab, + real_vocab, + temperature, + tp_group, + ): + logp, lse, hidden_2d, weight_c, target_1d, temp = _det_tp_run( + hidden, + weight, + bias, + target, + vocab_start=vocab_start, + global_vocab=global_vocab, + real_vocab=real_vocab, + temperature=temperature, + tp_group=tp_group, + ) + ctx.save_for_backward( + hidden_2d, + weight_c, + bias.contiguous() if bias is not None else hidden_2d.new_empty(0), + target_1d, + lse, + temp if temp is not None else hidden_2d.new_empty(0), + ) + ctx.has_bias = bias is not None + ctx.vocab_start = int(vocab_start) + ctx.real_vocab = int(real_vocab) + ctx.tp_group = tp_group + ctx.lead_shape = hidden.shape[:-1] + ctx.hidden_dtype = hidden.dtype + ctx.weight_dtype = weight.dtype + ctx.bias_dtype = bias.dtype if bias is not None else None + return logp.reshape(ctx.lead_shape), lse.reshape(ctx.lead_shape) + + @staticmethod + def backward(ctx, grad_logp, grad_lse): + if grad_logp is None and grad_lse is None: + return (None, None, None, None, None, None, None, None, None) + hidden_2d, weight, bias, target, lse, temp = ctx.saved_tensors + local_vocab = weight.size(0) + owned = (target >= ctx.vocab_start) & (target < ctx.vocab_start + local_vocab) + local_idx = (target - ctx.vocab_start).clamp(0, max(local_vocab - 1, 0)) + local_real_end = max(0, min(local_vocab, ctx.real_vocab - ctx.vocab_start)) + with torch.no_grad(): + dlogits = _det_reference_dlogits( + hidden_2d, + weight, + bias if ctx.has_bias else None, + local_idx, + owned, + lse, + temp if temp.numel() else None, + grad_logp, + grad_lse, + real_vocab_end=local_real_end, + ) + grad_hidden = grad_weight = grad_bias = None + if ctx.needs_input_grad[0]: + grad_hidden = dlogits.matmul(weight.float()) + dist = _require_distributed_initialized() + dist.all_reduce(grad_hidden, op=dist.ReduceOp.SUM, group=ctx.tp_group) + grad_hidden = grad_hidden.reshape((*tuple(ctx.lead_shape), weight.size(1))).to( + ctx.hidden_dtype + ) + if ctx.needs_input_grad[1]: + grad_weight = dlogits.transpose(0, 1).matmul(hidden_2d.float()).to(ctx.weight_dtype) + if ctx.has_bias and ctx.needs_input_grad[2]: + grad_bias = dlogits.sum(0).to(ctx.bias_dtype) + return grad_hidden, grad_weight, grad_bias, None, None, None, None, None, None + + +def triton_deterministic_linear_logp_tp( + hidden: torch.Tensor, + lm_head_weight_shard: torch.Tensor, + target_ids: torch.Tensor, + bias_shard: Optional[torch.Tensor] = None, + *, + vocab_start_index: int, + global_vocab_size: Optional[int] = None, + real_vocab_size: int = -1, + temperature: Optional[torch.Tensor] = None, + tp_group: Any = None, +): + """Strict tensor-parallel Triton selected logprob. + + Each rank runs the frozen-contract local kernels on its vocab shard and the + per-rank ``(lse, target-logit)`` stats merge through the shared explicit + ascending-rank chain, so the result is bitwise-stable for a fixed TP + topology and identical on every rank. The high-performance fused TP path is + unchanged.""" + if tp_group is None: + raise ValueError("strict TP linear_logp requires a TP process group") + if global_vocab_size is None: + dist = _require_distributed_initialized() + global_vocab_size = lm_head_weight_shard.size(0) * dist.get_world_size(tp_group) + real_vocab = int(global_vocab_size) if int(real_vocab_size) < 0 else int(real_vocab_size) + if torch.is_grad_enabled() and ( + hidden.requires_grad + or lm_head_weight_shard.requires_grad + or (bias_shard is not None and bias_shard.requires_grad) + ): + return _DetTritonTensorParallelLinearLogpAutograd.apply( + hidden, + lm_head_weight_shard, + bias_shard, + target_ids, + int(vocab_start_index), + int(global_vocab_size), + real_vocab, + temperature, + tp_group, + ) + logp, lse, _hidden, _weight, _target, _temp = _det_tp_run( + hidden, + lm_head_weight_shard, + bias_shard, + target_ids, + vocab_start=int(vocab_start_index), + global_vocab=int(global_vocab_size), + real_vocab=real_vocab, + temperature=temperature, + tp_group=tp_group, + ) + return logp.reshape(hidden.shape[:-1]), lse.reshape(hidden.shape[:-1]) diff --git a/tests/test_linear_logp.py b/tests/test_linear_logp.py index 1135bad8..72fcd497 100644 --- a/tests/test_linear_logp.py +++ b/tests/test_linear_logp.py @@ -1261,8 +1261,301 @@ def test_registry_dispatch_matches_native(): from rl_engine.platforms.device import device_ctx op = kernel_registry.get_op("linear_logp") - device = device_ctx.device if device_ctx.device_type == "cuda" else "cpu" + # ROCm reports device_type "rocm" but tensors still live on torch "cuda" + # devices; only true CPU hosts fall back to CPU inputs. + device = device_ctx.device if device_ctx.device_type in ("cuda", "rocm") else "cpu" hidden, weight, target, bias = _inputs(6, device=device) out = op(hidden, weight, target, bias) ref = NativeLinearLogpOp()(hidden, weight, target, bias) assert torch.allclose(out.cpu(), ref.cpu(), atol=1e-3) + + +# --------------------------------------------------------------------------- +# Deterministic (strict-contract) Triton linear_logp + + +def _det_triton_inputs(seed, *, n=96, d=192, v=4096, dtype=torch.bfloat16, bias=True): + generator = torch.Generator(device="cpu").manual_seed(seed) + hidden = (torch.randn(n, d, generator=generator) * 0.5).to("cuda", dtype) + weight = (torch.randn(v, d, generator=generator) * 0.05).to("cuda", dtype) + bias_t = (torch.randn(v, generator=generator) * 0.1).to("cuda", dtype) if bias else None + target = torch.randint(0, v, (n,), generator=generator).to("cuda") + return hidden, weight, target, bias_t + + +@requires_triton_cuda +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32], ids=["bf16", "fp32"]) +def test_triton_det_contract_is_bitwise_batch_invariant(dtype): + from rl_engine.kernels.ops.triton.loss.linear_logp import ( + TRITON_LINEAR_LOGP_CONTRACT_VERSION, + TRITON_N_SPLIT_CONTRACT, + triton_deterministic_linear_logp, + ) + + hidden, weight, target, bias = _det_triton_inputs(220, dtype=dtype) + real_vocab = weight.size(0) - 27 + target = target.remainder(real_vocab) + temperature = torch.linspace(0.7, 1.3, hidden.size(0), device="cuda", dtype=torch.float32) + + full_logp, full_lse = triton_deterministic_linear_logp( + hidden, weight, target, bias, temperature=temperature, real_vocab_size=real_vocab + ) + row = 37 + row_logp, row_lse = triton_deterministic_linear_logp( + hidden[row : row + 1], + weight, + target[row : row + 1], + bias, + temperature=temperature[row : row + 1], + real_vocab_size=real_vocab, + ) + again_logp, again_lse = triton_deterministic_linear_logp( + hidden, weight, target, bias, temperature=temperature, real_vocab_size=real_vocab + ) + + assert TRITON_LINEAR_LOGP_CONTRACT_VERSION == "triton-fused-linear-logp-contract-v1" + assert TRITON_N_SPLIT_CONTRACT == 64 + assert torch.equal(full_logp[row : row + 1].view(torch.int32), row_logp.view(torch.int32)) + assert torch.equal(full_lse[row : row + 1].view(torch.int32), row_lse.view(torch.int32)) + assert torch.equal(full_logp.view(torch.int32), again_logp.view(torch.int32)) + assert torch.equal(full_lse.view(torch.int32), again_lse.view(torch.int32)) + + +@requires_triton_cuda +def test_triton_det_contract_logits_return_and_padding_do_not_change_logp_bits(): + from rl_engine.kernels.ops.triton.loss.linear_logp import triton_deterministic_linear_logp + + hidden, weight, target, bias = _det_triton_inputs(221) + real_vocab = weight.size(0) - 27 + target = target.remainder(real_vocab) + temperature = torch.full((hidden.size(0),), 0.85, device="cuda", dtype=torch.float32) + + plain_logp, plain_lse = triton_deterministic_linear_logp( + hidden, weight, target, bias, temperature=temperature, real_vocab_size=real_vocab + ) + logits_logp, logits_lse, logits = triton_deterministic_linear_logp( + hidden, + weight, + target, + bias, + temperature=temperature, + return_logits=True, + real_vocab_size=real_vocab, + ) + + assert torch.equal(plain_logp.view(torch.int32), logits_logp.view(torch.int32)) + assert torch.equal(plain_lse.view(torch.int32), logits_lse.view(torch.int32)) + assert torch.isneginf(logits[:, real_vocab:]).all() + assert bool((plain_logp <= 0).all()) + + # Stored logits are unscaled: bias applied, temperature never applied. + reference_logits = torch.nn.functional.linear( + hidden.float(), weight[:real_vocab].float(), bias[:real_vocab].float() + ) + assert torch.allclose(logits[:, :real_vocab], reference_logits, atol=2e-2, rtol=2e-2) + reference_logp = torch.log_softmax(reference_logits / temperature[:, None], dim=-1) + reference_logp = reference_logp.gather(1, target[:, None]).squeeze(1) + assert torch.allclose(plain_logp, reference_logp, atol=2e-2, rtol=2e-2) + + +@requires_triton_cuda +@pytest.mark.parametrize("with_bias", [False, True], ids=["no-bias", "bias"]) +def test_triton_det_contract_backward_matches_logp_and_lse_reference(with_bias): + from rl_engine.kernels.ops.triton.loss.linear_logp import triton_deterministic_linear_logp + + hidden, weight, target, bias = _det_triton_inputs(222, bias=with_bias) + real_vocab = weight.size(0) - 27 + target = target.remainder(real_vocab) + n = hidden.size(0) + hidden = hidden.detach().requires_grad_(True) + weight = weight.detach().requires_grad_(True) + if with_bias: + bias = bias.detach().requires_grad_(True) + temperature = torch.full((n,), 0.9, device="cuda", dtype=torch.float32) + grad_logp = torch.linspace(-0.5, 0.5, n, device="cuda") + grad_lse = torch.linspace(0.25, -0.25, n, device="cuda") + + logp, lse = triton_deterministic_linear_logp( + hidden, weight, target, bias, temperature=temperature, real_vocab_size=real_vocab + ) + torch.autograd.backward((logp, lse), (grad_logp, grad_lse)) + + ref_hidden = hidden.detach().float().requires_grad_(True) + ref_weight = weight.detach().float().requires_grad_(True) + ref_bias = bias.detach().float().requires_grad_(True) if with_bias else None + ref_logits = torch.nn.functional.linear(ref_hidden, ref_weight, ref_bias) + ref_logits = ref_logits / temperature[:, None] + columns = torch.arange(ref_logits.size(1), device="cuda") + ref_logits = ref_logits.masked_fill(columns[None, :] >= real_vocab, float("-inf")) + ref_lse = torch.logsumexp(ref_logits, dim=-1) + ref_logp = ref_logits.gather(1, target[:, None]).squeeze(1) - ref_lse + torch.autograd.backward((ref_logp, ref_lse), (grad_logp, grad_lse)) + + assert torch.allclose(hidden.grad.float(), ref_hidden.grad, atol=2e-1, rtol=5e-2) + assert torch.allclose(weight.grad.float(), ref_weight.grad, atol=2e-1, rtol=5e-2) + if with_bias: + assert torch.allclose(bias.grad.float(), ref_bias.grad, atol=2e-2, rtol=5e-2) + + +@requires_triton_cuda +def test_triton_det_rejects_bad_inputs(): + from rl_engine.kernels.ops.triton.loss.linear_logp import triton_deterministic_linear_logp + + hidden, weight, target, bias = _det_triton_inputs(223, n=4, v=128) + with pytest.raises(ValueError, match="real_vocab_size"): + triton_deterministic_linear_logp(hidden, weight, target, bias, real_vocab_size=0) + with pytest.raises(ValueError, match="temperature"): + triton_deterministic_linear_logp( + hidden, weight, target, bias, temperature=torch.zeros(1, device="cuda") + ) + with pytest.raises(ValueError, match="real vocabulary"): + triton_deterministic_linear_logp( + hidden, weight, torch.full_like(target, weight.size(0)), bias + ) + + +def _det_triton_tp_worker(rank, world_size, init_method, result_queue): + try: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + torch.distributed.init_process_group( + backend="nccl", init_method=init_method, rank=rank, world_size=world_size + ) + from rl_engine.kernels.ops.triton.loss.linear_logp import ( + triton_deterministic_linear_logp, + triton_deterministic_linear_logp_tp, + ) + + generator = torch.Generator(device="cpu").manual_seed(224) + n, d, vocab = 48, 128, 2048 + real_vocab = vocab - 50 + hidden = (torch.randn(n, d, generator=generator) * 0.5).to(device, torch.bfloat16) + weight = (torch.randn(vocab, d, generator=generator) * 0.05).to(device, torch.bfloat16) + bias = (torch.randn(vocab, generator=generator) * 0.1).to(device, torch.bfloat16) + target = torch.randint(0, real_vocab, (n,), generator=generator).to(device) + temperature = torch.linspace(0.8, 1.2, n).to(device) + shard = vocab // world_size + vocab_start = rank * shard + + hidden_leaf = hidden.clone().requires_grad_(True) + weight_shard = weight[vocab_start : vocab_start + shard].contiguous().requires_grad_(True) + logp, lse = triton_deterministic_linear_logp_tp( + hidden_leaf, + weight_shard, + target, + bias[vocab_start : vocab_start + shard].contiguous(), + vocab_start_index=vocab_start, + global_vocab_size=vocab, + real_vocab_size=real_vocab, + temperature=temperature, + tp_group=torch.distributed.group.WORLD, + ) + (logp.sum() + 0.5 * lse.sum()).backward() + + again_logp, again_lse = triton_deterministic_linear_logp_tp( + hidden, + weight_shard.detach(), + target, + bias[vocab_start : vocab_start + shard].contiguous(), + vocab_start_index=vocab_start, + global_vocab_size=vocab, + real_vocab_size=real_vocab, + temperature=temperature, + tp_group=torch.distributed.group.WORLD, + ) + single_logp, single_lse = triton_deterministic_linear_logp( + hidden, weight, target, bias, temperature=temperature, real_vocab_size=real_vocab + ) + + ref_hidden = hidden.float().requires_grad_(True) + ref_weight = weight.float().requires_grad_(True) + ref_logits = torch.nn.functional.linear(ref_hidden, ref_weight, bias.float()) + ref_logits = ref_logits / temperature[:, None] + columns = torch.arange(vocab, device=device) + ref_logits = ref_logits.masked_fill(columns[None, :] >= real_vocab, float("-inf")) + ref_lse = torch.logsumexp(ref_logits, dim=-1) + ref_logp = ref_logits.gather(1, target[:, None]).squeeze(1) - ref_lse + (ref_logp.sum() + 0.5 * ref_lse.sum()).backward() + + result_queue.put( + { + "ok": True, + "rank": rank, + "repeat_bitwise": bool( + torch.equal(logp.detach().view(torch.int32), again_logp.view(torch.int32)) + and torch.equal(lse.detach().view(torch.int32), again_lse.view(torch.int32)) + ), + "close_to_reference": bool( + torch.allclose(logp.detach(), ref_logp.detach(), atol=2e-2, rtol=2e-2) + ), + "close_to_single_rank": bool( + torch.allclose(logp.detach(), single_logp, atol=1e-4, rtol=1e-5) + and torch.allclose(lse.detach(), single_lse, atol=1e-4, rtol=1e-5) + ), + "grad_hidden_close": bool( + torch.allclose(hidden_leaf.grad.float(), ref_hidden.grad, atol=2e-1, rtol=5e-2) + ), + "grad_weight_close": bool( + torch.allclose( + weight_shard.grad.float(), + ref_weight.grad[vocab_start : vocab_start + shard], + atol=2e-1, + rtol=5e-2, + ) + ), + "logp_bits": logp.detach().contiguous().view(torch.int32).cpu().tolist(), + } + ) + except Exception: # pragma: no cover - forwarded to the parent + result_queue.put({"ok": False, "rank": rank, "tb": traceback.format_exc()}) + raise + finally: + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +@requires_triton_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="strict TP test needs two GPUs") +@pytest.mark.parametrize( + "world_size", + [2, pytest.param(4, marks=pytest.mark.skipif(torch.cuda.device_count() < 4, reason="4 GPUs"))], +) +def test_triton_det_tp_matches_single_rank_and_replicates(world_size): + context = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmpdir: + init_method = (Path(tmpdir) / "det_tp_init").as_uri() + result_queue = context.Queue() + processes = [ + context.Process( + target=_det_triton_tp_worker, + args=(rank, world_size, init_method, result_queue), + ) + for rank in range(world_size) + ] + results = [] + try: + for process in processes: + process.start() + for _ in range(world_size): + try: + results.append(result_queue.get(timeout=300)) + except queue.Empty: + for process in processes: + if process.is_alive(): + process.terminate() + pytest.fail("timed out waiting for strict TP linear_logp workers") + finally: + for process in processes: + process.join(timeout=30) + if process.is_alive(): + process.terminate() + for result in results: + assert result["ok"], result.get("tb") + assert result["repeat_bitwise"] + assert result["close_to_reference"] + assert result["close_to_single_rank"] + assert result["grad_hidden_close"] + assert result["grad_weight_close"] + # Outputs are replicated: every rank holds identical bits. + for other in results[1:]: + assert results[0]["logp_bits"] == other["logp_bits"]