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: KotshinZ/gpt2-RMT-7
-datasets: HuggingFaceH4/Bespoke-Stratos-17k
 library_name: transformers
 model_name: gpt2-RMT-8
 tags:
 - generated_from_trainer
-- open-r1
 - trl
 - sft
 licence: license
@@ -13,7 +11,7 @@ licence: license
 
 # Model Card for gpt2-RMT-8
 
-This model is a fine-tuned version of [KotshinZ/gpt2-RMT-7](https://huggingface.co/KotshinZ/gpt2-RMT-7) on the [HuggingFaceH4/Bespoke-Stratos-17k](https://huggingface.co/datasets/HuggingFaceH4/Bespoke-Stratos-17k) dataset.
+This model is a fine-tuned version of [KotshinZ/gpt2-RMT-7](https://huggingface.co/KotshinZ/gpt2-RMT-7).
 It has been trained using [TRL](https://github.com/huggingface/trl).
 
 ## Quick start
@@ -29,14 +27,14 @@ 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/ch6o4apk) 
+[<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/4gmo4uif) 
 
 
 This model was trained with SFT.
 
 ### Framework versions
 
-- TRL: 0.16.0.dev0
+- TRL: 0.15.2
 - Transformers: 4.50.0.dev0
 - Pytorch: 2.5.1
 - Datasets: 3.3.2

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

@@ -1,13 +1,8 @@
 {
-    "eval_loss": 1.9101824760437012,
-    "eval_runtime": 2.9445,
-    "eval_samples": 100,
-    "eval_samples_per_second": 372.563,
-    "eval_steps_per_second": 93.395,
-    "total_flos": 2.3971976183808e+16,
-    "train_loss": 1.9497717069673088,
-    "train_runtime": 714.9694,
-    "train_samples": 16610,
-    "train_samples_per_second": 128.305,
-    "train_steps_per_second": 8.02
+    "total_flos": 5452245613150208.0,
+    "train_loss": 3.1222703313253013,
+    "train_runtime": 348.715,
+    "train_samples": 19883,
+    "train_samples_per_second": 7.605,
+    "train_steps_per_second": 0.238
 }

config.json

@@ -1,22 +1,126 @@
 {
-  "_name_or_path": "KotshinZ/gpt2-RMT-7",
   "activation_function": "gelu_new",
+  "align": "left",
   "architectures": [
-    "GPT2LMHeadModel"
+    "RecurrentMemoryTransformer"
   ],
   "attn_pdrop": 0.1,
+  "base_model_config": {
+    "_attn_implementation_autoset": true,
+    "_name_or_path": "openai-community/gpt2",
+    "activation_function": "gelu_new",
+    "add_cross_attention": false,
+    "architectures": [
+      "GPT2LMHeadModel"
+    ],
+    "attn_pdrop": 0.1,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 50256,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "embd_pdrop": 0.1,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 50256,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_epsilon": 1e-05,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "gpt2",
+    "n_ctx": 1024,
+    "n_embd": 768,
+    "n_head": 12,
+    "n_inner": null,
+    "n_layer": 12,
+    "n_positions": 1024,
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "reorder_and_upcast_attn": false,
+    "repetition_penalty": 1.0,
+    "resid_pdrop": 0.1,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "scale_attn_by_inverse_layer_idx": false,
+    "scale_attn_weights": true,
+    "sep_token_id": null,
+    "summary_activation": null,
+    "summary_first_dropout": 0.1,
+    "summary_proj_to_labels": true,
+    "summary_type": "cls_index",
+    "summary_use_proj": true,
+    "suppress_tokens": null,
+    "task_specific_params": {
+      "text-generation": {
+        "do_sample": true,
+        "max_length": 50
+      }
+    },
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": false,
+    "vocab_size": 50257
+  },
+  "base_model_type": "gpt2",
   "bos_token_id": 50256,
   "embd_pdrop": 0.1,
   "eos_token_id": 50256,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
   "initializer_range": 0.02,
+  "input_seg_len": 1004,
+  "is_memory_all": false,
   "layer_norm_epsilon": 1e-05,
-  "model_type": "gpt2",
+  "max_n_segments": 2,
+  "memory_size": 10,
+  "model_type": "rmt_gpt2",
   "n_ctx": 1024,
   "n_embd": 768,
   "n_head": 12,
   "n_inner": null,
   "n_layer": 12,
   "n_positions": 1024,
+  "num_mem_tokens": 10,
+  "output_seg_len": 1004,
   "reorder_and_upcast_attn": false,
   "resid_pdrop": 0.1,
   "scale_attn_by_inverse_layer_idx": false,
@@ -34,6 +138,6 @@
   },
   "torch_dtype": "bfloat16",
   "transformers_version": "4.50.0.dev0",
-  "use_cache": true,
+  "use_cache": false,
   "vocab_size": 50257
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:785b3bc458e117f5a369b848617750c5e7a0951dee5013c3a213521f521a682e
-size 326089656
+oid sha256:8a608b7f1a18a1d7c70b018ba505eb4edce00cd300ac3027ffc3d60408e9c2c8
+size 248915448

tokenizer_config.json

@@ -11,7 +11,6 @@
     }
   },
   "bos_token": "<|endoftext|>",
-  "chat_template": "{% for message in messages %}{{ message['role'] }}: {{ message['content'] }}\n{% endfor %}",
   "clean_up_tokenization_spaces": false,
   "eos_token": "<|endoftext|>",
   "extra_special_tokens": {},

train_results.json

@@ -1,8 +1,8 @@
 {
-    "total_flos": 2.3971976183808e+16,
-    "train_loss": 1.9497717069673088,
-    "train_runtime": 714.9694,
-    "train_samples": 16610,
-    "train_samples_per_second": 128.305,
-    "train_steps_per_second": 8.02
+    "total_flos": 5452245613150208.0,
+    "train_loss": 3.1222703313253013,
+    "train_runtime": 348.715,
+    "train_samples": 19883,
+    "train_samples_per_second": 7.605,
+    "train_steps_per_second": 0.238
 }

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4b104884d6ba1f77b5b69bc7b531f0c6d732ce72aed65b9dab9261c66a5de002
+oid sha256:ec32a102e3bf57cfe8a1fabd9e01b5ba70a29d6a8e8a4c20779e360df3b2ef1b
 size 7352