Upload folder using huggingface_hub

beagle.json

@@ -0,0 +1,99 @@
+{
+  "dataset.debug": false,
+  "dataset.git_diff": "",
+  "dataset.git_sha1": "unknown",
+  "dataset.manual_sample_ids": [],
+  "dataset.max_read_items": null,
+  "dataset.output_dir": "output",
+  "dataset.path": "/mnt/hg_cache/temp_llama_dataset/datasets/ds_Llama-2-7b-chat-hf",
+  "dataset.read_eagle_format": false,
+  "dataset.run_name": "temp_run",
+  "dataset.seed": 42,
+  "dataset_generation.batch_size": 1,
+  "dataset_generation.debug": false,
+  "dataset_generation.debug_target": null,
+  "dataset_generation.ds_prefix": "ds_",
+  "dataset_generation.git_diff": "",
+  "dataset_generation.git_sha1": "unknown",
+  "dataset_generation.max_length": 4096,
+  "dataset_generation.output_dir": "output",
+  "dataset_generation.run_name": "temp_run",
+  "dataset_generation.save_every": 1000,
+  "dataset_generation.seed": 42,
+  "inference.debug": false,
+  "inference.detail_time_stats": false,
+  "inference.draft_tree_shape": "mc_sim_7b_64",
+  "inference.git_diff": "",
+  "inference.git_sha1": "unknown",
+  "inference.interactive": false,
+  "inference.max_new_tokens": 512,
+  "inference.mode": "speculative",
+  "inference.output_dir": "output",
+  "inference.run_name": "temp_run",
+  "inference.seed": 42,
+  "modeling.add_noise": true,
+  "modeling.attention_offset": "random.randrange(0, 3)",
+  "modeling.attention_wind": "5",
+  "modeling.ckpt_path": null,
+  "modeling.debug": false,
+  "modeling.decoder_key_remap": {},
+  "modeling.draft_growing": false,
+  "modeling.dtype": "torch.float32",
+  "modeling.frozen_targets": [],
+  "modeling.git_diff": "",
+  "modeling.git_sha1": "unknown",
+  "modeling.layer_path": "model.layers",
+  "modeling.lmhead_path": "lm_head",
+  "modeling.load_config_from_model_path": false,
+  "modeling.model_path": "beagle/models/llama2/Llama-2-7b-chat-hf/",
+  "modeling.norm_path": "model.norm",
+  "modeling.output_dir": "output",
+  "modeling.reuse_layer": null,
+  "modeling.rotary_path": "model.rotary_emb",
+  "modeling.run_name": "temp_run",
+  "modeling.save_loading": true,
+  "modeling.seed": 42,
+  "modeling.strictly_follow_eagle_decoder": false,
+  "modeling.tokenizer_path": "meta-llama/Llama-2-7b-chat-hf",
+  "modeling.use_fc_eagle": false,
+  "modeling.use_lower_layers": 0,
+  "modeling.use_state_distill": false,
+  "training.adam_beta1": 0.9,
+  "training.adam_beta2": 0.95,
+  "training.bf16": true,
+  "training.ddp_find_unused_parameters": false,
+  "training.debug": false,
+  "training.eval_steps": 100,
+  "training.eval_strategy": "steps",
+  "training.filter_out_shorts": false,
+  "training.git_diff": "diff --git a/README.md b/README.md\nindex 906761d..da3f70e 100644\n--- a/README.md\n+++ b/README.md\n@@ -21,13 +21,14 @@ Training:\n ```sh\n wandb login\n # multi-GPU training:\n-torchrun --standalone --nnodes=1 --nproc-per-node=2 -m beagle.train \\\n+CUDA_VISIBLE_DEVICES=0,1,2,3 \\\n+    torchrun --standalone --nnodes=1 --nproc-per-node=4 -m beagle.train \\\n     --@llama2_7b_chat --@rtx4070tis_bs4_ctx4096 --modeling.use_lower_layers 0 \\\n     --dataset.path /mnt/hg_cache/temp_llama_dataset/datasets/ds_Llama-2-7b-chat-hf \\\n     --modeling.save_loading --training.report_to wandb\n \n # customized training:\n-CUDA_VISIBLE_DEVICES=1,2 \\\n+CUDA_VISIBLE_DEVICES=0,1 \\\n     torchrun --standalone --nnodes=1 --nproc-per-node=2 -m beagle.train \\\n     [email protected]_2.4B_instr --@rtx4070tis_dev_bs8 \\\n     --dataset.path /mnt/truenas_sync/beagle_train_data/datasets/ds_EXAONE-3.5-2.4B-Instruct \\\ndiff --git a/beagle/train.py b/beagle/train.py\nindex 74795e1..c25096e 100644\n--- a/beagle/train.py\n+++ b/beagle/train.py\n@@ -199,8 +199,6 @@ def train(configs, run_name, tokenizer, model,\n     if not configs.overwrite_output_dir:\n         assert not os.path.exists(configs.output_dir), configs.output_dir\n \n-    print('[rank]', rank, '/', world_size)\n-\n     random.seed(configs.seed)\n     torch.manual_seed(configs.seed)\n \n@@ -326,6 +324,9 @@ def main(config_file='beagle/configs.ini', **injects):\n     configs.set_obj('training.git_diff', git_diff)\n     master_print('[git]', git_sha1, '\\n', git_diff)\n \n+    print('[rank]', rank, '/', world_size)\n+    configs.set_obj('training.world_size', world_size)\n+\n     if rank == 0 and configs.training.report_to == 'wandb':\n         import wandb\n         wandb.init(",
+  "training.git_sha1": "453d65d9cbee2ac0dedb03c3c49d99b1344a81ed",
+  "training.gradient_accumulation_steps": 8,
+  "training.learning_rate": 3e-05,
+  "training.logging_steps": 5,
+  "training.lr_scheduler_type": "constant_with_warmup",
+  "training.max_grad_norm": 0.5,
+  "training.max_length": 4096,
+  "training.max_steps": -1,
+  "training.model_init_ckpt": null,
+  "training.num_train_epochs": 10,
+  "training.optim": "adamw_torch_fused",
+  "training.output_dir": "output",
+  "training.overwrite_output_dir": true,
+  "training.per_device_eval_batch_size": 1,
+  "training.per_device_train_batch_size": 2,
+  "training.project": "beagle",
+  "training.report_to": "wandb",
+  "training.resume_from_checkpoint": false,
+  "training.resume_wandb_runid": null,
+  "training.run_name": "temp_run",
+  "training.save_steps": 500,
+  "training.save_strategy": "steps",
+  "training.save_total_limit": 2,
+  "training.save_vram": true,
+  "training.seed": 42,
+  "training.tf32": false,
+  "training.use_eagle_pipeline": false,
+  "training.warmup_steps": 2000,
+  "training.world_size": 1
+}

config.json

@@ -0,0 +1,61 @@
+{
+  "_name_or_path": "meta-llama/Llama-2-7b-chat-hf",
+  "architectures": [
+    "LlamaForSpeculativeCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModelForSpeculativeCausalLM": "modeling_speculative_llama.LlamaForSpeculativeCausalLM"
+  },
+  "beagle_add_noise": true,
+  "beagle_attention_offset": "random.randrange(0, 3)",
+  "beagle_attention_wind": "5",
+  "beagle_ckpt_path": null,
+  "beagle_debug": false,
+  "beagle_decoder_key_remap": {},
+  "beagle_draft_growing": false,
+  "beagle_dtype": "torch.float32",
+  "beagle_frozen_targets": [],
+  "beagle_git_diff": "",
+  "beagle_git_sha1": "unknown",
+  "beagle_layer_path": "model.layers",
+  "beagle_lmhead_path": "lm_head",
+  "beagle_load_config_from_model_path": false,
+  "beagle_model_path": "beagle/models/llama2/Llama-2-7b-chat-hf/",
+  "beagle_norm_path": "model.norm",
+  "beagle_output_dir": "output",
+  "beagle_reuse_layer": null,
+  "beagle_rotary_path": "model.rotary_emb",
+  "beagle_run_name": "temp_run",
+  "beagle_save_loading": true,
+  "beagle_seed": 42,
+  "beagle_strictly_follow_eagle_decoder": false,
+  "beagle_tokenizer_path": "meta-llama/Llama-2-7b-chat-hf",
+  "beagle_use_fc_eagle": false,
+  "beagle_use_lower_layers": 0,
+  "beagle_use_state_distill": false,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 11008,
+  "max_position_embeddings": 4096,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 0,
+  "num_key_value_heads": 32,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "speculative_base_model": "meta-llama/Llama-2-7b-chat-hf",
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.47.1",
+  "use_cache": true,
+  "vocab_size": 32000
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "bos_token_id": 1,
+  "do_sample": true,
+  "eos_token_id": 2,
+  "max_length": 4096,
+  "pad_token_id": 0,
+  "temperature": 0.6,
+  "top_p": 0.9,
+  "transformers_version": "4.47.1"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b84ac6ec585166ca8ae8639bb5481bcfcc32fb2b443c96ebe0d1ebf828996df
+size 1858127224

modeling_speculative_llama.py

@@ -0,0 +1,88 @@
+from transformers.models.llama.modeling_llama import *
+from beagle.mixin import *
+
+
+class LlamaBeagleAttention(LlamaAttention, BeagleAttentionMixin):
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+
+        bsz, q_len, _ = hidden_states.size()
+        query_states, key_states, value_states = self.qkv_transform(
+            hidden_states, past_key_value, use_cache, position_embeddings, **kwargs)
+
+        ################################################
+        ### everything kept original starting from here
+        ################################################
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attention_mask is not None:  # no matter the length, we just slice it
+            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+            attn_weights = attn_weights + causal_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class LlamaBeagleDecoderLayer(LlamaDecoderLayer):
+    def __init__(self, config, layer_id):
+        super().__init__(config, layer_id)
+
+        if not config.beagle_use_fc_eagle:
+            delattr(self, 'self_attn')
+            recycle_vram()
+            self.self_attn = LlamaBeagleAttention(
+                config=config, layer_idx=layer_id
+            )
+
+        if config.beagle_strictly_follow_eagle_decoder:
+            delattr(self, 'input_layernorm')
+            recycle_vram()
+            self.input_layernorm = nn.Identity()
+
+
+class LlamaForSpeculativeCausalLM(LlamaForCausalLM, BeagleMixin):
+    _no_split_modules = ["LlamaDecoderLayer", "LlamaBeagleDecoderLayer"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        BeagleMixin.__init__(self, config)
+        self.speculative_decoder = LlamaBeagleDecoderLayer(config, layer_id=0)
+
+        self.post_init()
+
+    def forward(self, *args, **kwargs) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPast]:
+        return self.beagle_forward(*args, **kwargs)

special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "</s>",
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,44 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% for message in loop_messages %}{% if loop.index0 == 0 and system_message != false %}{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ bos_token + '[INST] ' + content.strip() + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{% generation %}{{ ' '  + content.strip() + ' ' + eos_token }}{% endgeneration %}{% endif %}{% endfor %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "legacy": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "</s>",
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}