From 18d5eb2d441e4340d9700f4007628a221f0ac145 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 21:40:06 -0400 Subject: [PATCH] fix: improve PEFT device consistency and XReg output slicing - Initialize LoRA parameters on the same device as the base linear layer - Load adapter weights directly to the model device instead of hardcoded CPU - Slice XReg linear regression outputs to match the specified sequence lengths - Replace batch-wide covariate normalization with per-input normalization in create_covariate_matrix to prevent cross-input scale leakage. - Refactor BatchedInContextXRegLinear.fit to solve ridge regression per instance rather than as a single global matrix solve, avoiding cross-contamination between batched inputs. - Truncate JAX regression outputs to the actual train/test lengths after the padded matrix multiply, fixing shape mismatches for non-power-of-2 horizons (e.g. horizon=24 was returning 32 elements). --- peft/adapters.py | 8 +++++--- src/timesfm/utils/xreg_lib.py | 4 ++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/peft/adapters.py b/peft/adapters.py index 951cfe1..acc6eee 100644 --- a/peft/adapters.py +++ b/peft/adapters.py @@ -110,9 +110,10 @@ class DoRALinear(nn.Module): in_f = base_linear.in_features out_f = base_linear.out_features + dev = base_linear.weight.device - self.lora_A = nn.Parameter(torch.empty(in_f, rank)) - self.lora_B = nn.Parameter(torch.zeros(rank, out_f)) + self.lora_A = nn.Parameter(torch.empty(in_f, rank, device=dev)) + self.lora_B = nn.Parameter(torch.zeros(rank, out_f, device=dev)) nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) # Magnitude vector — initialised from pretrained column norms. @@ -272,7 +273,8 @@ def load_adapter_weights(model: nn.Module, path: str) -> None: The model must already have adapters injected (via ``inject_adapters``) before calling this function. """ - tensors = load_file(path, device="cpu") + device = str(next(model.parameters()).device) + tensors = load_file(path, device=device) trainable = {n for n, p in model.named_parameters() if p.requires_grad} missing = trainable - set(tensors.keys()) if missing: diff --git a/src/timesfm/utils/xreg_lib.py b/src/timesfm/utils/xreg_lib.py index 4355261..2759b67 100644 --- a/src/timesfm/utils/xreg_lib.py +++ b/src/timesfm/utils/xreg_lib.py @@ -509,9 +509,9 @@ class BatchedInContextXRegLinear(BatchedInContextXRegBase): @ x_tr_j.T @ y_tr_j ) - outputs.append(np.array(x_te_j @ beta_hat)) + outputs.append(np.array(x_te_j @ beta_hat)[:tel]) if debug_info: - outputs_context.append(np.array(x_tr_raw_j @ beta_hat)) + outputs_context.append(np.array(x_tr_raw_j @ beta_hat)[:trl]) train_idx += trl test_idx += tel