school assignment 4

test_basic.infini.py

@@ -4,7 +4,7 @@
 
 import torch
 from torch.nn import functional as F
-from transformers import AutoTokenizer, pipeline
+from transformers import AutoTokenizer
 from infini_gemma import GemmaForCausalLM
 from infini_gemma import GemmaConfig
 
@@ -24,95 +24,71 @@
 
 # Create the Gemma model with Infini-attention
 model = GemmaForCausalLM(config)
-# model = model.from_pretrained("google/gemma-2b")
-pretrained_model = GemmaForCausalLM.from_pretrained("google/gemma-2b")
-# Step 4: Transfer weights
-# Note: This is a simplified example; you need to ensure that each parameter's dimensions match.
-for param in model.named_parameters():
-    name = param[0]
-    if name in pretrained_model.state_dict():
-        # Check if dimensions match, and only then assign the weights
-        if param[1].size() == pretrained_model.state_dict()[name].size():
-            param[1].data = pretrained_model.state_dict()[name].data.clone()
-        else:
-            print(f"Skipping {name} due to size mismatch.")
+
+#  pretrained_model 초기화와 가중치 이식 코드 통합 및 개선
+with torch.no_grad():
+    pretrained_model = GemmaForCausalLM.from_pretrained("google/gemma-2b")
+    for name, param in model.named_parameters():
+        if name in pretrained_model.state_dict():
+            pretrained_param = pretrained_model.state_dict()[name]
+            if param.size() == pretrained_param.size():
+                param.copy_(pretrained_param)
+            else:
+                print(f"Skipping {name} due to size mismatch.")
+
 print(model)
 model.to(device)
 
-# Generate some dummy input data
 tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
 text = """This work introduces an efficient method to scale Transformer-based"""
 longtext = """The new memory states M s and z s are then passed to the next segment S + 1, building in a recurrence in each attention layer. The right side term σ (K ) T V in Eq. (4) is known as an associative binding operator (Smolensky, 1990; Hebb, 2005; Schlag et al., 2020).
 Inspired by the success of delta rule (Munkhdalai et al., 2019; Schlag et al., 2020; 2021), we have also incorporated it into our Infini-attention. The delta rule attempts a slightly improved memory update by first retrieving existing value entries and subtracting them from the new values before applying the associative bindings as new update."""
 
-encoded = tokenizer(
-    text,
-    return_tensors="pt",
-)
-# attention_mask = torch.ones_like(input_ids)
+encoded = tokenizer(text, return_tensors="pt")
 encoded["labels"] = encoded["input_ids"].clone()
 
-long_encoded = tokenizer(
-    longtext,
-    return_tensors="pt",
-)
-# attention_mask = torch.ones_like(input_ids)
+long_encoded = tokenizer(longtext, return_tensors="pt")
 long_encoded["labels"] = long_encoded["input_ids"].clone()
 
 print(encoded)
-# Test the forward pass
-outputs = model(**encoded.to(device))  # position_ids=position_ids)
+
+outputs = model(**encoded.to(device))
 print("Short Text Loss")
 print(outputs.loss)
-outputs.loss.backward()  # Test the backward pass
+outputs.loss.backward()
 
-outputs = model(**long_encoded.to(device))  # position_ids=position_ids)
+outputs = model(**long_encoded.to(device))
 print("Long Text Loss")
 print(outputs.loss)
-outputs.loss.backward()  # Test the backward pass
+outputs.loss.backward()
 
 print("backprop done")
 
-
-# Step 1: Get effective batch size and sequence length
-batch_size = encoded["input_ids"].shape[0]
-sequence_length = encoded["input_ids"].shape[1]
-
-# Step 2: Prepare input data for generation
 input_ids = encoded["input_ids"]
 attention_mask = encoded.get("attention_mask", None)
-
-# Step 3: Initialize past
 past = None
 
-# Step 4: Start generation loop
-for _ in range(10):  # 10 is the number of new tokens to generate
-    with torch.no_grad():
-        # Get next token scores
+#  torch.inference_mode() 사용으로 불필요한 그래프 추적 방지
+for _ in range(10):
+    with torch.inference_mode():
         outputs = model(
-            input_ids,
+            input_ids=input_ids,
             attention_mask=attention_mask,
             use_cache=True,
             past_key_values=past,
         )
         next_token_logits = outputs.logits[:, -1, :]
         past = outputs.past_key_values
 
-        # Perform sampling to get the next token
-        next_token = torch.multinomial(
-            F.softmax(next_token_logits, dim=-1), num_samples=1
-        )
+        next_token = torch.multinomial(F.softmax(next_token_logits, dim=-1), num_samples=1)
 
-        # Update input_ids, attention_mask, and past
         input_ids = torch.cat([input_ids, next_token], dim=-1)
         if attention_mask is not None:
             attention_mask = F.pad(attention_mask, (0, 1), value=1)
 
-# Step 5: Return generated sequence
 generated_sequence = tokenizer.decode(input_ids[0], skip_special_tokens=False)
 print("generated_sequence:", generated_sequence)
 
-# Test .generate() method
 generated = model.generate(
     **encoded,
     max_new_tokens=32,