Model save

PreTrainedRMTConfig.py

@@ -0,0 +1,56 @@
+import os
+import json
+from transformers import PretrainedConfig
+
+class PreTrainedRMTConfig(PretrainedConfig):
+    """
+    Recurrent Memory Transformer の設定クラス
+    """
+    
+    model_type = "rmt"
+    
+    # マッピング情報を追加（設定クラスとモデルクラスの関連付け）
+    auto_map = {
+        "AutoModelForCausalLM": "open_r1.rmt.RecurrentMemoryTransofomer.RecurrentMemoryTransformer"
+    }
+    
+    def __init__(
+        self,
+        base_model_config=None,
+        is_memory_all=True,
+        max_n_segments=1,
+        input_seg_len=512,
+        output_seg_len=512,
+        align="left",
+        num_mem_tokens=10,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.base_model_config = base_model_config
+        self.is_memory_all = is_memory_all
+        self.max_n_segments = max_n_segments
+        self.input_seg_len = input_seg_len
+        self.output_seg_len = output_seg_len
+        self.align = align
+        self.num_mem_tokens = num_mem_tokens
+        
+        if base_model_config is not None:
+            if type(base_model_config) is not dict:
+                dict_config: dict = base_model_config.to_dict()
+            else:
+                dict_config: dict = base_model_config
+                
+            for key, value in dict_config.items():
+                setattr(self, key, value)
+            self.base_model_type = dict_config.get("model_type")
+            if self.base_model_type is None:
+                raise ValueError("base_model_configにmodel_typeが指定されていません。")
+            PreTrainedRMTConfig.model_type = "rmt_" + self.base_model_type   
+    """
+    def __repr__(self):
+        return f"PreTrainedRMTConfig(is_memory_all={self.is_memory_all}, max_n_segments={self.max_n_segments}, " \
+               f"input_seg_len={self.input_seg_len}, output_seg_len={self.output_seg_len}, " \
+               f"align='{self.align}', num_mem_tokens={self.num_mem_tokens})"
+    """
+
+PreTrainedRMTConfig.register_for_auto_class()

README.md

@@ -1,11 +1,9 @@
 ---
 base_model: openai-community/gpt2
-datasets: HuggingFaceFW/fineweb-edu
 library_name: transformers
 model_name: gpt2-RMT-2
 tags:
 - generated_from_trainer
-- open-r1
 - trl
 - sft
 licence: license
@@ -13,7 +11,7 @@ licence: license
 
 # Model Card for gpt2-RMT-2
 
-This model is a fine-tuned version of [openai-community/gpt2](https://huggingface.co/openai-community/gpt2) on the [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) dataset.
+This model is a fine-tuned version of [openai-community/gpt2](https://huggingface.co/openai-community/gpt2).
 It has been trained using [TRL](https://github.com/huggingface/trl).
 
 ## Quick start
@@ -29,7 +27,7 @@ print(output["generated_text"])
 
 ## Training procedure
 
-[<img src="https://raw.githubusercontent.com/wandb/assets/main/wandb-github-badge-28.svg" alt="Visualize in Weights & Biases" width="150" height="24"/>](https://wandb.ai/s18574s18574-/huggingface/runs/r46ahkkj) 
+[<img src="https://raw.githubusercontent.com/wandb/assets/main/wandb-github-badge-28.svg" alt="Visualize in Weights & Biases" width="150" height="24"/>](https://wandb.ai/shin2021001-osaka-city-university/huggingface/runs/5xyxy6y1) 
 
 
 This model was trained with SFT.

RecurrentMemoryTransofomer.py

@@ -0,0 +1,171 @@
+import torch
+from transformers import PreTrainedModel, AutoModelForCausalLM, AutoConfig
+from transformers.models.auto.auto_factory import _BaseAutoModelClass
+from open_r1.rmt.MemoryCell import MemoryCell
+from open_r1.rmt.RecurrentWrapper import RecurrentWrapper
+from open_r1.rmt.PreTrainedRMTConfig import PreTrainedRMTConfig
+
+
+# @register_for_auto_class("AutoModelForCausalLM")
+class RecurrentMemoryTransformer(PreTrainedModel):
+    """
+    Recurrent Memory Transformer モデルクラス
+    長い文脈をセグメント単位で処理し、メモリを使って情報を保持するトランスフォーマーモデル
+    """
+    
+    config_class = PreTrainedRMTConfig
+    auto_model_class = "AutoModelForCausalLM"
+    
+    # マッピングを定義してAutoクラスが適切なモデルを見つけられるようにする
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+    
+    # AUTO_MAPを定義（モデル名からクラスへのマッピング）
+    AUTO_MAP = {
+        "AutoModelForCausalLM": "RecurrentMemoryTransformer",
+    }
+    
+    def __init__(self, config, base_model=None):
+        """
+        初期化
+        
+        Parameters
+        ----------
+        config : PreTrainedRMTConfig
+            モデルの設定
+        base_model : PreTrainedModel, optional
+            ベースとなるトランスフォーマーモデル
+        """
+        super().__init__(config)
+        
+        # base_modelが指定されていない場合は、configから自動生成
+        if base_model is None:
+            # ベースモデルのタイプを確認
+            if not hasattr(config, "base_model_type"):
+                raise ValueError("configにbase_model_typeが指定されていません。RMTの設定にはベースモデルタイプが必要です。")
+            base_model_type = config.base_model_type
+            
+            # ベースモデル用の設定を作成
+            base_config = AutoConfig.from_pretrained(base_model_type)
+            
+            # RMT固有のパラメータを除外してベースモデルの設定を作成
+            rmt_specific_params = ['model_type', 'is_memory_all', 'max_n_segments', 'input_seg_len', 
+                                  'output_seg_len', 'align', 'num_mem_tokens', 'base_model_type']
+            for key, value in config.__dict__.items():
+                if key not in rmt_specific_params and not key.startswith('_'):
+                    setattr(base_config, key, value)
+            
+            # ベースモデルを作成
+            base_model = AutoModelForCausalLM.from_config(base_config)
+        
+        # MemoryCellとRecurrentWrapperの初期化
+        memory_cell = MemoryCell(base_model, config.num_mem_tokens)
+        self.recurrent_wrapper = RecurrentWrapper(
+            memory_cell=memory_cell,
+            is_memory_all=config.is_memory_all,
+            max_n_segments=config.max_n_segments,
+            input_seg_len=config.input_seg_len,
+            output_seg_len=config.output_seg_len,
+            align=config.align
+        )
+        
+    def get_base_model(self):
+        """
+        ベースモデルを取得
+        """
+        return self.recurrent_wrapper.memory_cell.model
+    
+    def forward(self, input_ids=None, attention_mask=None, labels=None, labels_mask=None, 
+                inputs_embeds=None, output_attentions=None, output_hidden_states=None):
+        """
+        モデルの順伝播
+        
+        Parameters
+        ----------
+        input_ids : torch.Tensor, optional
+            入力テンソル
+        attention_mask : torch.Tensor, optional
+            アテンションマスク
+        labels : torch.Tensor, optional
+            ラベルテンソル
+        labels_mask : torch.Tensor, optional
+            ラベルマスク
+        inputs_embeds : torch.Tensor, optional
+            入力埋め込み
+        output_attentions : bool, optional
+            アテンション重みを出力するかどうか
+        output_hidden_states : bool, optional
+            隠れ状態を出力するかどうか
+        """
+        forward_kwargs = {}
+        if input_ids is not None:
+            forward_kwargs["input_ids"] = input_ids
+        if labels is not None:
+            forward_kwargs["labels"] = labels
+        if attention_mask is not None:
+            forward_kwargs["attention_mask"] = attention_mask
+        if labels_mask is not None:
+            forward_kwargs["labels_mask"] = labels_mask
+        if inputs_embeds is not None:
+            forward_kwargs["inputs_embeds"] = inputs_embeds
+        if output_attentions is not None:
+            forward_kwargs["output_attentions"] = output_attentions
+        if output_hidden_states is not None:
+            forward_kwargs["output_hidden_states"] = output_hidden_states
+        
+        #forward_kwargs.update(kwargs)
+        
+        # 通常の順伝播処理
+        out = self.recurrent_wrapper.forward(**forward_kwargs)
+        """
+        # デバッグ出力を削除（または必要に応じてコメント化）
+        # print(out["loss"])
+        
+        # 分散環境で損失が二重計算されないよう、ワールドサイズで割る
+        # これは処理済みの場合は不要なので、環境変数などで制御することも可能
+        if torch.distributed.is_initialized() and "loss" in out and out["loss"] is not None:
+            # 既にDeepSpeedが処理している可能性があるため、確認が必要
+            # テスト目的で一時的に追加（実際の環境に合わせて調整が必要）
+            # world_size = torch.distributed.get_world_size()
+            # out["loss"] = out["loss"] / world_size
+            pass
+        """
+        return out
+    
+    def generate(self, **kwargs):
+        """
+        テキスト生成
+        """
+        return self.recurrent_wrapper.generate(**kwargs)
+    
+    def generate_with_tokenizer(self, tokenizer, input_text, **kwargs):
+        """
+        トークナイザーを用いたテキスト生成
+        """
+        return self.recurrent_wrapper.generate_with_tokenizer(tokenizer, input_text, **kwargs)
+    
+    def get_input_embeddings(self):
+        """
+        入力埋め込みを取得
+        """
+        return self.get_base_model().get_input_embeddings()
+    
+    def set_input_embeddings(self, embeddings):
+        """
+        入力埋め込みを設定
+        """
+        self.get_base_model().set_input_embeddings(embeddings)
+    
+    def get_output_embeddings(self):
+        """
+        出力埋め込みを取得
+        """
+        return self.get_base_model().get_output_embeddings()
+    
+    def resize_token_embeddings(self, new_num_tokens):
+        """
+        トークン埋め込みのサイズを変更
+        """
+        self.get_base_model().resize_token_embeddings(new_num_tokens)
+        return self.get_input_embeddings()
+
+RecurrentMemoryTransformer.register_for_auto_class("AutoModelForCausalLM")

all_results.json

@@ -3,10 +3,10 @@
     "eval_samples": 100,
     "eval_samples_per_second": 193.544,
     "eval_steps_per_second": 12.097,
-    "total_flos": 541940076511232.0,
-    "train_loss": 11.141571969696969,
-    "train_runtime": 94.2109,
+    "total_flos": 5419400756723712.0,
+    "train_loss": 3.615518465909091,
+    "train_runtime": 337.9584,
     "train_samples": 19883,
-    "train_samples_per_second": 11.221,
-    "train_steps_per_second": 0.35
+    "train_samples_per_second": 31.279,
+    "train_steps_per_second": 0.976
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f139c841af94b44aea936b1787cdb6832e0a84c3740b43cac34d0196953244ed
+oid sha256:203c58315e462136da93aea0cc81387cd50fe8ac5e3d6dca237c751d223ab4fd
 size 248915448

train_results.json

@@ -1,8 +1,8 @@
 {
-    "total_flos": 541940076511232.0,
-    "train_loss": 11.141571969696969,
-    "train_runtime": 94.2109,
+    "total_flos": 5419400756723712.0,
+    "train_loss": 3.615518465909091,
+    "train_runtime": 337.9584,
     "train_samples": 19883,
-    "train_samples_per_second": 11.221,
-    "train_steps_per_second": 0.35
+    "train_samples_per_second": 31.279,
+    "train_steps_per_second": 0.976
 }

trainer_state.json

@@ -1,50 +1,310 @@
 {
   "best_metric": null,
   "best_model_checkpoint": null,
-  "epoch": 0.09984871406959153,
+  "epoch": 0.9984871406959153,
   "eval_steps": 100,
-  "global_step": 33,
+  "global_step": 330,
   "is_hyper_param_search": false,
   "is_local_process_zero": true,
   "is_world_process_zero": true,
   "log_history": [
     {
       "epoch": 0.030257186081694403,
-      "grad_norm": 39.661752788828096,
-      "learning_rate": 1.796093065705644e-05,
-      "loss": 13.7688,
-      "mean_token_accuracy": 0.14352177157998086,
+      "grad_norm": 48.18317253284181,
+      "learning_rate": 6.060606060606061e-06,
+      "loss": 7.4023,
+      "mean_token_accuracy": 0.13559286370873452,
       "step": 10
     },
     {
       "epoch": 0.060514372163388806,
-      "grad_norm": 21.083752762804057,
-      "learning_rate": 8.382180034472353e-06,
-      "loss": 10.7016,
-      "mean_token_accuracy": 0.18965651467442513,
+      "grad_norm": 25.86717975049105,
+      "learning_rate": 1.2121212121212122e-05,
+      "loss": 6.8148,
+      "mean_token_accuracy": 0.14477298334240912,
       "step": 20
     },
     {
       "epoch": 0.0907715582450832,
-      "grad_norm": 16.256136796751306,
-      "learning_rate": 5.234682881719766e-07,
-      "loss": 9.5031,
-      "mean_token_accuracy": 0.23370618522167205,
+      "grad_norm": 9.978340525429944,
+      "learning_rate": 1.8181818181818182e-05,
+      "loss": 5.6102,
+      "mean_token_accuracy": 0.17600756287574768,
       "step": 30
     },
     {
-      "epoch": 0.09984871406959153,
-      "mean_token_accuracy": 0.23802571247021356,
-      "step": 33,
-      "total_flos": 541940076511232.0,
-      "train_loss": 11.141571969696969,
-      "train_runtime": 94.2109,
-      "train_samples_per_second": 11.221,
-      "train_steps_per_second": 0.35
+      "epoch": 0.12102874432677761,
+      "grad_norm": 4.525649385270933,
+      "learning_rate": 1.9972599751485225e-05,
+      "loss": 4.6477,
+      "mean_token_accuracy": 0.2445007838308811,
+      "step": 40
+    },
+    {
+      "epoch": 0.15128593040847202,
+      "grad_norm": 3.175409164823092,
+      "learning_rate": 1.9838755799290993e-05,
+      "loss": 4.0312,
+      "mean_token_accuracy": 0.30081592202186586,
+      "step": 50
+    },
+    {
+      "epoch": 0.1815431164901664,
+      "grad_norm": 2.010787736128858,
+      "learning_rate": 1.9594929736144978e-05,
+      "loss": 3.6578,
+      "mean_token_accuracy": 0.33934821784496305,
+      "step": 60
+    },
+    {
+      "epoch": 0.2118003025718608,
+      "grad_norm": 1.1782901971905857,
+      "learning_rate": 1.9243847161266924e-05,
+      "loss": 3.3934,
+      "mean_token_accuracy": 0.371051961183548,
+      "step": 70
+    },
+    {
+      "epoch": 0.24205748865355523,
+      "grad_norm": 0.9889630793327328,
+      "learning_rate": 1.8789432636206197e-05,
+      "loss": 3.3324,
+      "mean_token_accuracy": 0.378506475687027,
+      "step": 80
+    },
+    {
+      "epoch": 0.2723146747352496,
+      "grad_norm": 0.8299178994463533,
+      "learning_rate": 1.8236765814298328e-05,
+      "loss": 3.2738,
+      "mean_token_accuracy": 0.38666853606700896,
+      "step": 90
+    },
+    {
+      "epoch": 0.30257186081694404,
+      "grad_norm": 0.8704351270978057,
+      "learning_rate": 1.7592024657977432e-05,
+      "loss": 3.282,
+      "mean_token_accuracy": 0.38509487807750703,
+      "step": 100
+    },
+    {
+      "epoch": 0.30257186081694404,
+      "eval_runtime": 0.225,
+      "eval_samples_per_second": 213.326,
+      "eval_steps_per_second": 13.333,
+      "step": 100
+    },
+    {
+      "epoch": 0.3328290468986384,
+      "grad_norm": 0.7957814410722591,
+      "learning_rate": 1.686241637868734e-05,
+      "loss": 3.2516,
+      "mean_token_accuracy": 0.39028556644916534,
+      "step": 110
+    },
+    {
+      "epoch": 0.3630862329803328,
+      "grad_norm": 0.7305776403219497,
+      "learning_rate": 1.6056096871376667e-05,
+      "loss": 3.2766,
+      "mean_token_accuracy": 0.38682268708944323,
+      "step": 120
+    },
+    {
+      "epoch": 0.39334341906202724,
+      "grad_norm": 0.7140945050608877,
+      "learning_rate": 1.5182079544175957e-05,
+      "loss": 3.25,
+      "mean_token_accuracy": 0.39047552198171614,
+      "step": 130
+    },
+    {
+      "epoch": 0.4236006051437216,
+      "grad_norm": 0.7281100436036338,
+      "learning_rate": 1.4250134562400301e-05,
+      "loss": 3.2254,
+      "mean_token_accuracy": 0.39294967502355577,
+      "step": 140
+    },
+    {
+      "epoch": 0.45385779122541603,
+      "grad_norm": 0.7678124673120136,
+      "learning_rate": 1.3270679633174219e-05,
+      "loss": 3.2156,
+      "mean_token_accuracy": 0.39464686065912247,
+      "step": 150
+    },
+    {
+      "epoch": 0.48411497730711045,
+      "grad_norm": 0.7134325713906087,
+      "learning_rate": 1.2254663551538047e-05,
+      "loss": 3.1895,
+      "mean_token_accuracy": 0.39698781073093414,
+      "step": 160
+    },
+    {
+      "epoch": 0.5143721633888049,
+      "grad_norm": 0.723513310530251,
+      "learning_rate": 1.121344380981082e-05,
+      "loss": 3.2254,
+      "mean_token_accuracy": 0.39353512674570085,
+      "step": 170
+    },
+    {
+      "epoch": 0.5446293494704992,
+      "grad_norm": 0.7185772760297996,
+      "learning_rate": 1.015865963834808e-05,
+      "loss": 3.1977,
+      "mean_token_accuracy": 0.39701825082302095,
+      "step": 180
+    },
+    {
+      "epoch": 0.5748865355521936,
+      "grad_norm": 0.6834580433094664,
+      "learning_rate": 9.102101896903084e-06,
+      "loss": 3.2086,
+      "mean_token_accuracy": 0.39654194712638857,
+      "step": 190
+    },
+    {
+      "epoch": 0.6051437216338881,
+      "grad_norm": 0.7155622236506424,
+      "learning_rate": 8.055581271005292e-06,
+      "loss": 3.2094,
+      "mean_token_accuracy": 0.3948154032230377,
+      "step": 200
+    },
+    {
+      "epoch": 0.6051437216338881,
+      "eval_runtime": 0.2213,
+      "eval_samples_per_second": 216.915,
+      "eval_steps_per_second": 13.557,
+      "step": 200
+    },
+    {
+      "epoch": 0.6354009077155824,
+      "grad_norm": 0.6809195871075517,
+      "learning_rate": 7.0307962467172555e-06,
+      "loss": 3.1988,
+      "mean_token_accuracy": 0.3979633778333664,
+      "step": 210
+    },
+    {
+      "epoch": 0.6656580937972768,
+      "grad_norm": 0.707315821198318,
+      "learning_rate": 6.039202339608432e-06,
+      "loss": 3.202,
+      "mean_token_accuracy": 0.3954067021608353,
+      "step": 220
+    },
+    {
+      "epoch": 0.6959152798789713,
+      "grad_norm": 0.6904494791556431,
+      "learning_rate": 5.091884039764321e-06,
+      "loss": 3.1965,
+      "mean_token_accuracy": 0.39610737562179565,
+      "step": 230
+    },
+    {
+      "epoch": 0.7261724659606656,
+      "grad_norm": 0.6699258512413571,
+      "learning_rate": 4.19943090428802e-06,
+      "loss": 3.2105,
+      "mean_token_accuracy": 0.39350710064172745,
+      "step": 240
+    },
+    {
+      "epoch": 0.75642965204236,
+      "grad_norm": 0.74188228093062,
+      "learning_rate": 3.37181918238904e-06,
+      "loss": 3.2074,
+      "mean_token_accuracy": 0.39521674066782,
+      "step": 250
+    },
+    {
+      "epoch": 0.7866868381240545,
+      "grad_norm": 0.6872378478922624,
+      "learning_rate": 2.618300296308135e-06,
+      "loss": 3.2215,
+      "mean_token_accuracy": 0.3938901349902153,
+      "step": 260
+    },
+    {
+      "epoch": 0.8169440242057489,
+      "grad_norm": 0.7136188203818471,
+      "learning_rate": 1.947297424689414e-06,
+      "loss": 3.1988,
+      "mean_token_accuracy": 0.39579751938581464,
+      "step": 270
+    },
+    {
+      "epoch": 0.8472012102874432,
+      "grad_norm": 0.6657712543261404,
+      "learning_rate": 1.3663113444380905e-06,
+      "loss": 3.202,
+      "mean_token_accuracy": 0.3932719826698303,
+      "step": 280
+    },
+    {
+      "epoch": 0.8774583963691377,
+      "grad_norm": 0.72900722083696,
+      "learning_rate": 8.818365836066101e-07,
+      "loss": 3.1625,
+      "mean_token_accuracy": 0.4012471452355385,
+      "step": 290
+    },
+    {
+      "epoch": 0.9077155824508321,
+      "grad_norm": 0.6725834304571868,
+      "learning_rate": 4.992888225905467e-07,
+      "loss": 3.2246,
+      "mean_token_accuracy": 0.3945300817489624,
+      "step": 300
+    },
+    {
+      "epoch": 0.9077155824508321,
+      "eval_runtime": 0.2212,
+      "eval_samples_per_second": 217.045,
+      "eval_steps_per_second": 13.565,
+      "step": 300
+    },
+    {
+      "epoch": 0.9379727685325264,
+      "grad_norm": 0.7015653789891463,
+      "learning_rate": 2.2294435517691504e-07,
+      "loss": 3.1926,
+      "mean_token_accuracy": 0.3973464578390121,
+      "step": 310
+    },
+    {
+      "epoch": 0.9682299546142209,
+      "grad_norm": 0.6744365913859031,
+      "learning_rate": 5.5892286176932875e-08,
+      "loss": 3.1965,
+      "mean_token_accuracy": 0.39635206162929537,
+      "step": 320
+    },
+    {
+      "epoch": 0.9984871406959153,
+      "grad_norm": 0.6617082977205768,
+      "learning_rate": 0.0,
+      "loss": 3.2031,
+      "mean_token_accuracy": 0.3948866441845894,
+      "step": 330
+    },
+    {
+      "epoch": 0.9984871406959153,
+      "step": 330,
+      "total_flos": 5419400756723712.0,
+      "train_loss": 3.615518465909091,
+      "train_runtime": 337.9584,
+      "train_samples_per_second": 31.279,
+      "train_steps_per_second": 0.976
     }
   ],
   "logging_steps": 10,
-  "max_steps": 33,
+  "max_steps": 330,
   "num_input_tokens_seen": 0,
   "num_train_epochs": 1,
   "save_steps": 500,
@@ -60,7 +320,7 @@
       "attributes": {}
     }
   },
-  "total_flos": 541940076511232.0,
+  "total_flos": 5419400756723712.0,
   "train_batch_size": 8,
   "trial_name": null,
   "trial_params": null

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9e730128fa97641e51393b5facdd56eec2c7c21fa8e4909ab7c389ef5e29c0f5
+oid sha256:88757a2b8b1c49ba8b8f97161dacdb95ca6b6fd7a56d93df7f7c1df11e238bd7
 size 7352