|
36 | 36 | approx_topk_threshold |
37 | 37 | ) |
38 | 38 |
|
39 | | -from src.modeling_utils import ( |
40 | | - FastLoRAProjection, BaseModelOutputWithPastAndPredictorLoss |
41 | | -) |
42 | | - |
43 | 39 | logger = logging.get_logger(__name__) |
44 | 40 |
|
| 41 | +@dataclass |
| 42 | +class BaseModelOutputWithPastAndPredictorLoss(ModelOutput): |
| 43 | + loss: Optional[torch.FloatTensor] = None |
| 44 | + last_hidden_state: Optional[torch.FloatTensor] = None |
| 45 | + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None |
| 46 | + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None |
| 47 | + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None |
| 48 | + |
| 49 | + |
| 50 | +class FastLoRAProjection(nn.Module): |
| 51 | + def __init__(self, hidden_size, intermediate_size, lora_size): |
| 52 | + super().__init__() |
| 53 | + self.hidden_size = hidden_size |
| 54 | + self.intermediate_size = intermediate_size |
| 55 | + self.lora_size = lora_size |
| 56 | + # Force creation of linear layers with actual tensors (not meta tensors) |
| 57 | + self.down = nn.Linear(hidden_size, lora_size, bias=False) |
| 58 | + self.up = nn.Linear(lora_size, intermediate_size, bias=False) |
| 59 | + # Pre-allocate buffers on CPU initially |
| 60 | + self.register_buffer('intermediate', torch.zeros(1, lora_size)) |
| 61 | + self.register_buffer('output', torch.zeros(1, intermediate_size)) |
| 62 | + |
| 63 | + def to(self, *args, **kwargs): |
| 64 | + # Move buffers to same device as model when .to() is called |
| 65 | + device = args[0] if args else kwargs.get('device') |
| 66 | + |
| 67 | + if device: |
| 68 | + self.intermediate = self.intermediate.to(device) |
| 69 | + self.output = self.output.to(device) |
| 70 | + return super().to(*args, **kwargs) |
| 71 | + |
| 72 | + def _fix_unloaded_weights(self): |
| 73 | + out = self.to_empty(device="cpu") |
| 74 | + with torch.no_grad(): |
| 75 | + torch.nn.init.xavier_normal_(out.down.weight) |
| 76 | + torch.nn.init.zeros_(out.up.weight) # Initialize up projection to zeros for stable training |
| 77 | + return out |
| 78 | + |
| 79 | + def _resize_buffers(self, batch_size: int, dtype: torch.dtype): |
| 80 | + if self.intermediate.size(0) != batch_size: |
| 81 | + self.intermediate.resize_(batch_size, self.lora_size) |
| 82 | + self.intermediate = self.intermediate.to(dtype=dtype) |
| 83 | + self.intermediate.fill_(0.0) # Explicitly initialize with zeros |
| 84 | + self.output.resize_(batch_size, self.intermediate_size) |
| 85 | + self.output = self.output.to(dtype=dtype) |
| 86 | + self.output.fill_(0.0) # Explicitly initialize with zeros |
| 87 | + |
| 88 | + def forward(self, x): |
| 89 | + batch_size = x.size(0) |
| 90 | + |
| 91 | + # Check if gradients are required (training mode) |
| 92 | + if self.training: |
| 93 | + # Use regular matrix multiplication for gradient computation |
| 94 | + intermediate = torch.mm(x, self.down.weight.t()) |
| 95 | + output = torch.mm(intermediate, self.up.weight.t()) |
| 96 | + return output |
| 97 | + else: |
| 98 | + # # Use optimized in-place operations for inference |
| 99 | + # intermediate = torch.mm(x, self.down.weight.t()) |
| 100 | + # output = torch.mm(intermediate, self.up.weight.t()) |
| 101 | + # return output |
| 102 | + |
| 103 | + self._resize_buffers(batch_size, x.dtype) |
| 104 | + torch.mm(x, self.down.weight.t(), out=self.intermediate) |
| 105 | + torch.mm(self.intermediate, self.up.weight.t(), out=self.output) |
| 106 | + return self.output |
45 | 107 |
|
46 | 108 | class SkipMLP(nn.Module): |
47 | 109 | def __init__(self, hidden_size: int, intermediate_size: int, sparsity: float, bias: bool = False): |
@@ -415,6 +477,15 @@ def __init__(self, config): |
415 | 477 |
|
416 | 478 | # Initialize weights and apply final processing |
417 | 479 | self.post_init() |
| 480 | + |
| 481 | + @classmethod |
| 482 | + def from_pretrained(cls, *args, **kwargs): |
| 483 | + out = super(SkipConnectionModelForCausalLM, cls).from_pretrained(*args, **kwargs) |
| 484 | + for module in out.modules(): |
| 485 | + if any(hasattr(p, 'is_meta') and p.is_meta for p in module.parameters()) and \ |
| 486 | + hasattr(module, '_fix_unloaded_weights'): |
| 487 | + module = module._fix_unloaded_weights() |
| 488 | + return out |
418 | 489 |
|
419 | 490 | def get_input_embeddings(self): |
420 | 491 | return self.model.embed_tokens |
|
0 commit comments