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

@@ -0,0 +1,58 @@
+---
+base_model: KotshinZ/gpt2-RMT-7-mem512
+library_name: transformers
+model_name: gpt2-RMT-8-mem512
+tags:
+- generated_from_trainer
+- trl
+- sft
+licence: license
+---
+
+# Model Card for gpt2-RMT-8-mem512
+
+This model is a fine-tuned version of [KotshinZ/gpt2-RMT-7-mem512](https://huggingface.co/KotshinZ/gpt2-RMT-7-mem512).
+It has been trained using [TRL](https://github.com/huggingface/trl).
+
+## Quick start
+
+```python
+from transformers import pipeline
+
+question = "If you had a time machine, but could only go to the past or the future once and never return, which would you choose and why?"
+generator = pipeline("text-generation", model="KotshinZ/gpt2-RMT-8-mem512", device="cuda")
+output = generator([{"role": "user", "content": question}], max_new_tokens=128, return_full_text=False)[0]
+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/shin2021001-osaka-city-university/huggingface/runs/zgmj6bcp) 
+
+
+This model was trained with SFT.
+
+### Framework versions
+
+- TRL: 0.15.2
+- Transformers: 4.50.0.dev0
+- Pytorch: 2.5.1
+- Datasets: 3.3.2
+- Tokenizers: 0.21.0
+
+## Citations
+
+
+
+Cite TRL as:
+    
+```bibtex
+@misc{vonwerra2022trl,
+	title        = {{TRL: Transformer Reinforcement Learning}},
+	author       = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallouédec},
+	year         = 2020,
+	journal      = {GitHub repository},
+	publisher    = {GitHub},
+	howpublished = {\url{https://github.com/huggingface/trl}}
+}
+```

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

@@ -0,0 +1,8 @@
+{
+    "total_flos": 5426859755962368.0,
+    "train_loss": 3.1510416666666665,
+    "train_runtime": 387.2622,
+    "train_samples": 19883,
+    "train_samples_per_second": 13.394,
+    "train_steps_per_second": 0.418
+}

config.json

@@ -0,0 +1,143 @@
+{
+  "activation_function": "gelu_new",
+  "align": "left",
+  "architectures": [
+    "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": 512,
+  "is_memory_all": false,
+  "layer_norm_epsilon": 1e-05,
+  "max_n_segments": 2,
+  "memory_size": 512,
+  "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": 512,
+  "reorder_and_upcast_attn": false,
+  "resid_pdrop": 0.1,
+  "scale_attn_by_inverse_layer_idx": false,
+  "scale_attn_weights": true,
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "task_specific_params": {
+    "text-generation": {
+      "do_sample": true,
+      "max_length": 50
+    }
+  },
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.50.0.dev0",
+  "use_cache": false,
+  "vocab_size": 50257
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:43aa3154a3a0507399f2b63925dbc848879161df546b22ae0e149c43b6ed5434
+size 248915448

special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|endoftext|>",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,21 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

train_results.json

@@ -0,0 +1,8 @@
+{
+    "total_flos": 5426859755962368.0,
+    "train_loss": 3.1510416666666665,
+    "train_runtime": 387.2622,
+    "train_samples": 19883,
+    "train_samples_per_second": 13.394,
+    "train_steps_per_second": 0.418
+}

trainer_state.json

@@ -0,0 +1,178 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.9969230769230769,
+  "eval_steps": 100,
+  "global_step": 162,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.06153846153846154,
+      "grad_norm": 0.39951827848719723,
+      "learning_rate": 1.1764705882352942e-05,
+      "loss": 3.1582,
+      "mean_token_accuracy": 0.400910259783268,
+      "step": 10
+    },
+    {
+      "epoch": 0.12307692307692308,
+      "grad_norm": 0.3995616499609791,
+      "learning_rate": 1.9978883431348845e-05,
+      "loss": 3.1668,
+      "mean_token_accuracy": 0.39937195032835004,
+      "step": 20
+    },
+    {
+      "epoch": 0.18461538461538463,
+      "grad_norm": 0.4102168990934247,
+      "learning_rate": 1.9605953553832987e-05,
+      "loss": 3.1387,
+      "mean_token_accuracy": 0.4016988441348076,
+      "step": 30
+    },
+    {
+      "epoch": 0.24615384615384617,
+      "grad_norm": 0.40391230141561146,
+      "learning_rate": 1.8783859964390466e-05,
+      "loss": 3.1551,
+      "mean_token_accuracy": 0.3992693811655045,
+      "step": 40
+    },
+    {
+      "epoch": 0.3076923076923077,
+      "grad_norm": 0.3961653949779693,
+      "learning_rate": 1.755104284557221e-05,
+      "loss": 3.1457,
+      "mean_token_accuracy": 0.40172072798013686,
+      "step": 50
+    },
+    {
+      "epoch": 0.36923076923076925,
+      "grad_norm": 0.41939623098641426,
+      "learning_rate": 1.5965147355676344e-05,
+      "loss": 3.1492,
+      "mean_token_accuracy": 0.4015924736857414,
+      "step": 60
+    },
+    {
+      "epoch": 0.4307692307692308,
+      "grad_norm": 0.3945808975159374,
+      "learning_rate": 1.4100328205214161e-05,
+      "loss": 3.1387,
+      "mean_token_accuracy": 0.4018391355872154,
+      "step": 70
+    },
+    {
+      "epoch": 0.49230769230769234,
+      "grad_norm": 0.37735963614306256,
+      "learning_rate": 1.204378226506365e-05,
+      "loss": 3.1488,
+      "mean_token_accuracy": 0.4004943951964378,
+      "step": 80
+    },
+    {
+      "epoch": 0.5538461538461539,
+      "grad_norm": 0.4114554888279876,
+      "learning_rate": 9.891671337699603e-06,
+      "loss": 3.1559,
+      "mean_token_accuracy": 0.4003040686249733,
+      "step": 90
+    },
+    {
+      "epoch": 0.6153846153846154,
+      "grad_norm": 0.41103454513473137,
+      "learning_rate": 7.74462573818606e-06,
+      "loss": 3.1543,
+      "mean_token_accuracy": 0.40160409063100816,
+      "step": 100
+    },
+    {
+      "epoch": 0.6153846153846154,
+      "eval_runtime": 0.4543,
+      "eval_samples_per_second": 52.823,
+      "eval_steps_per_second": 4.402,
+      "step": 100
+    },
+    {
+      "epoch": 0.676923076923077,
+      "grad_norm": 0.38569131648987776,
+      "learning_rate": 5.7030389324864845e-06,
+      "loss": 3.1824,
+      "mean_token_accuracy": 0.397557806968689,
+      "step": 110
+    },
+    {
+      "epoch": 0.7384615384615385,
+      "grad_norm": 0.4255553103146562,
+      "learning_rate": 3.862373250574626e-06,
+      "loss": 3.1828,
+      "mean_token_accuracy": 0.3972348183393478,
+      "step": 120
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 0.432279104045337,
+      "learning_rate": 2.308696173983711e-06,
+      "loss": 3.1387,
+      "mean_token_accuracy": 0.40221917182207106,
+      "step": 130
+    },
+    {
+      "epoch": 0.8615384615384616,
+      "grad_norm": 0.4143971545965478,
+      "learning_rate": 1.1146559160270875e-06,
+      "loss": 3.1297,
+      "mean_token_accuracy": 0.4037084937095642,
+      "step": 140
+    },
+    {
+      "epoch": 0.9230769230769231,
+      "grad_norm": 0.41227809966322443,
+      "learning_rate": 3.360844720863765e-07,
+      "loss": 3.1383,
+      "mean_token_accuracy": 0.40221751779317855,
+      "step": 150
+    },
+    {
+      "epoch": 0.9846153846153847,
+      "grad_norm": 0.4212135535860477,
+      "learning_rate": 9.38697756023288e-09,
+      "loss": 3.1223,
+      "mean_token_accuracy": 0.4041751459240913,
+      "step": 160
+    },
+    {
+      "epoch": 0.9969230769230769,
+      "mean_token_accuracy": 0.3931814655661583,
+      "step": 162,
+      "total_flos": 5426859755962368.0,
+      "train_loss": 3.1510416666666665,
+      "train_runtime": 387.2622,
+      "train_samples_per_second": 13.394,
+      "train_steps_per_second": 0.418
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 162,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": false,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 5426859755962368.0,
+  "train_batch_size": 8,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64aa7be91907887688fd3f87d1956ad836561b1853152b26a33b9fff2b072cbd
+size 7352