File size: 3,054 Bytes
199efb6 | 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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
try:
from .configuration_oddevendumb import OddEvenDumbConfig
except ImportError:
from configuration_oddevendumb import OddEvenDumbConfig
class BinarizeSTE(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
return torch.where(input >= 0.0, 1.0, -1.0)
@staticmethod
def backward(ctx, grad_output):
return grad_output
def binarize(tensor):
return BinarizeSTE.apply(tensor)
class OddEvenDumbPreTrainedModel(PreTrainedModel):
config_class = OddEvenDumbConfig
base_model_prefix = "oddevendumb"
def _init_weights(self, module):
pass
class OddEvenDumbForCausalLM(OddEvenDumbPreTrainedModel):
def __init__(self, config: OddEvenDumbConfig):
super().__init__(config)
self.config = config
# 1ビット動作を保証するための明示的なパラメータ定義
self.emb_weight = nn.Parameter(torch.randn(config.vocab_size, config.embed_dim) * 0.02)
self.w_ih = nn.Parameter(torch.randn(config.hidden_dim, config.embed_dim) * 0.02)
self.w_hh = nn.Parameter(torch.randn(config.hidden_dim, config.hidden_dim) * 0.02)
self.fc_weight = nn.Parameter(torch.randn(config.vocab_size, config.hidden_dim) * 0.02)
self.post_init()
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
labels: Optional[torch.LongTensor] = None,
return_dict: Optional[bool] = None,
**kwargs,
) -> Union[Tuple, CausalLMOutputWithPast]:
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
bin_emb = binarize(self.emb_weight)
bin_w_ih = binarize(self.w_ih)
bin_w_hh = binarize(self.w_hh)
bin_fc = binarize(self.fc_weight)
batch_size, seq_len = input_ids.shape
embeds = torch.nn.functional.embedding(input_ids, bin_emb)
h_t = torch.zeros(batch_size, self.config.hidden_dim, dtype=embeds.dtype, device=embeds.device)
hidden_states = []
for t in range(seq_len):
x_t = embeds[:, t, :]
h_t = torch.tanh(
torch.matmul(x_t, bin_w_ih.t()) + torch.matmul(h_t, bin_w_hh.t())
)
hidden_states.append(h_t.unsqueeze(1))
out = torch.cat(hidden_states, dim=1)
logits = torch.matmul(out, bin_fc.t())
loss = None
if labels is not None:
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
loss_fct = nn.CrossEntropyLoss()
loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
if not return_dict:
output = (logits,)
return ((loss,) + output) if loss is not None else output
return CausalLMOutputWithPast(loss=loss, logits=logits)
|