Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+wandb/debug-internal.log filter=lfs diff=lfs merge=lfs -text
+wandb/run-20240416_172306-uygw9yfk/run-uygw9yfk.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20240416_205309-xdytsc71/logs/debug-internal.log filter=lfs diff=lfs merge=lfs -text
+wandb/run-20240416_205309-xdytsc71/run-xdytsc71.wandb filter=lfs diff=lfs merge=lfs -text

accelerate_config.yaml

@@ -0,0 +1,17 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+gpu_ids: all
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

checkpoint-50000-epoch-0/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6501c2c0a4b20fadfdb42633a77b04e780b08a701b1043be50ef4fa8029eb4c5
+size 6283286904

checkpoint-50000-epoch-0/model_1.safetensors

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

checkpoint-50000-epoch-0/optimizer.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b4b0fa6a4b48ec9e9817f31f7956f46662000678470e2b911c0aaf0c97da8a0
+size 12566610474

checkpoint-50000-epoch-0/random_states_0.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9aca503cc09e63ca033e29a437a20cc580a9c1db27fef2174e533f58ba275879
+size 16100

checkpoint-50000-epoch-0/random_states_1.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31831c2134536b1e81ba1e763e72b2ff98a14a83774fcfb30d153a66dca7879c
+size 16100

checkpoint-50000-epoch-0/random_states_2.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a628258539b4090ce50e9faf5fda4d613f523ca957f3e837c02d316e4b20122
+size 16100

checkpoint-50000-epoch-0/random_states_3.pkl

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

checkpoint-50000-epoch-0/random_states_4.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28ca4240374ff4b93ad0537aca2f28bfc293153a29ee8069cf09d088ca30fee7
+size 16100

checkpoint-50000-epoch-0/random_states_5.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d6f3577977e8c32eac49b1c5136c6718fcd9c66051b703ba6e305cca03a8fb0
+size 16100

checkpoint-50000-epoch-0/random_states_6.pkl

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

checkpoint-50000-epoch-0/random_states_7.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08d860c07ef8d57c8162394106fcd87c34e7924d859b28b4b292e9e792a96af2
+size 16100

checkpoint-50000-epoch-0/scheduler.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61d45b6fd8dfc8830f2ac6212faa2f2fe74d12e2ca5a21e1d98c8dea13c837b7
+size 1064

config.json

@@ -0,0 +1,27 @@
+{
+  "_name_or_path": "sanchit-gandhi/Mistral-7B-v0.1-6-layer",
+  "architectures": [
+    "MistralForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 14336,
+  "max_position_embeddings": 32768,
+  "model_type": "mistral",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 6,
+  "num_key_value_heads": 8,
+  "output_router_logits": true,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "sliding_window": 4096,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.0.dev0",
+  "use_cache": true,
+  "vocab_size": 32000
+}

config_dummy.yaml

@@ -0,0 +1,30 @@
+# Model arguments
+model_name_or_path: sanchit-gandhi/tiny-random-MistralForCausalLM-1-layer
+teacher_model_name_or_path: sanchit-gandhi/tiny-random-MistralForCausalLM-1-layer
+dtype: bfloat16
+load_teacher_in_4bit: true
+
+# Data arguments
+train_dataset_name: HuggingFaceTB/cosmopedia-100k
+train_dataset_config_name: default
+train_split_name: train[:100]
+eval_split_name: train[-10:]
+num_train_epochs: 2
+
+# Training arguments
+do_train: true
+do_eval: true
+per_device_eval_batch_size: 4
+per_device_train_batch_size: 4
+learning_rate: 0.0003
+warmup_steps: 10
+gradient_checkpointing: true
+dataloader_num_workers: 2
+preprocessing_num_workers: 2
+save_strategy: epoch
+evaluation_strategy: epoch
+logging_steps: 25
+overwrite_output_dir: true
+output_router_logits: true
+report_to: wandb
+output_dir: ./

config_mistral.yaml

@@ -0,0 +1,42 @@
+# Model arguments
+model_name_or_path: sanchit-gandhi/Mistral-7B-v0.1-6-layer
+teacher_model_name_or_path: mistralai/Mistral-7B-v0.1
+dtype: bfloat16
+load_teacher_in_4bit: true
+
+# Data arguments
+train_dataset_name: HuggingFaceTB/cosmopedia
+train_dataset_config_name:
+  - auto_math_text
+  - khanacademy
+  - openstax
+  - stanford
+  - stories
+  - web_samples_v1
+  - web_samples_v2
+  - wikihow
+train_split_name: train[1000:]
+eval_split_name: train[:1000]
+max_steps: 50000
+
+# Training arguments
+do_train: true
+do_eval: true
+per_device_eval_batch_size: 8
+per_device_train_batch_size: 8
+learning_rate: 0.0003
+warmup_steps: 500
+gradient_checkpointing: true
+dataloader_num_workers: 4
+preprocessing_num_workers: 32
+ddp_timeout: 7200
+save_strategy: steps
+save_steps: 5000
+evaluation_strategy: steps
+eval_steps: 5000
+logging_steps: 25
+output_router_logits: true
+report_to: wandb
+output_dir: ./
+overwrite_output_dir: true
+save_total_limit: 1

config_mistral_100k.yaml

@@ -0,0 +1,31 @@
+# Model arguments
+model_name_or_path: sanchit-gandhi/Mistral-7B-v0.1-6-layer
+teacher_model_name_or_path: mistralai/Mistral-7B-v0.1
+dtype: bfloat16
+load_teacher_in_4bit: true
+
+# Data arguments
+train_dataset_name: HuggingFaceTB/cosmopedia-100k
+train_dataset_config_name: default
+train_split_name: train[:-1000]
+eval_split_name: train[-1000:]
+num_train_epochs: 10
+
+# Training arguments
+do_train: true
+do_eval: true
+per_device_eval_batch_size: 8
+per_device_train_batch_size: 8
+learning_rate: 0.0003
+warmup_steps: 500
+gradient_checkpointing: true
+dataloader_num_workers: 4
+preprocessing_num_workers: 32
+ddp_timeout: 7200
+save_strategy: epoch
+evaluation_strategy: epoch
+logging_steps: 25
+overwrite_output_dir: true
+output_router_logits: true
+report_to: wandb
+output_dir: ./

dummy_config.yaml

@@ -0,0 +1,17 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+gpu_ids: all
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 1
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "max_length": 2048,
+  "transformers_version": "4.40.0.dev0"
+}

model-00001-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2680ce25539a3c26efbd0f16858e75f000cd77a2b5e5b9a779fa1c9372f99fd1
+size 4987196936

model-00002-of-00002.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,64 @@
+{
+  "metadata": {
+    "total_size": 6283280384
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00002-of-00002.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.norm.weight": "model-00002-of-00002.safetensors"
+  }
+}

run_distillation.py

@@ -0,0 +1,1476 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Training langauge models Whisper model for conditional language modelling tasks via teacher-student distillation.
+"""
+# You can also adapt this script for your own distillation tasks. Pointers for this are left as comments.
+
+import logging
+import math
+import os
+import re
+import shutil
+import sys
+import time
+from dataclasses import dataclass, field
+from functools import partial
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import datasets
+import numpy as np
+import torch
+import torch.nn as nn
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from datasets import (
+    Dataset,
+    DatasetDict,
+    IterableDataset,
+    IterableDatasetDict,
+    concatenate_datasets,
+    interleave_datasets,
+    load_dataset,
+)
+from huggingface_hub import create_repo, get_full_repo_name, upload_folder
+from peft import LoraConfig, get_peft_model
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BatchEncoding,
+    BitsAndBytesConfig,
+    HfArgumentParser,
+    PreTrainedTokenizerBase,
+    Seq2SeqTrainingArguments,
+    get_scheduler,
+    set_seed,
+)
+from transformers.utils import check_min_version
+from transformers.utils.versions import require_version
+
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version("4.34.0.dev0")
+
+require_version("datasets>=2.14.6", "To fix: `pip install --upgrade datasets`")
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to distill from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained Whisper model or model identifier from huggingface.co/models"}
+    )
+    teacher_model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained teacher model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "Pretrained config name or path if not the same as model_name"},
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"},
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    subfolder: str = field(
+        default="",
+        metadata={
+            "help": "In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can"
+            "specify the folder name here."
+        },
+    )
+    token: str = field(
+        default=None,
+        metadata={
+            "help": (
+                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
+                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
+            )
+        },
+    )
+    attn_implementation: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Which attention implementation to use in the encoder and decoder attention layers. Can be one of:\n"
+                "1. `eager` or `None`: default Transformers attention implementation.\n"
+                "2. `sdpa`: Flash Attention through PyTorch SDPA. Requires `torch>=2.1`. Recommended for hardware where Flash Attention 2 is not supported, e.g. Turing GPUs, (T4, RTX 2080).\n"
+                "3. `flash_attn_2`: Flash Attention 2 through the Flash Attention package https://github.com/Dao-AILab/flash-attention. **Always** recommended on supported hardware (Ampere, Ada, or Hopper GPUs, e.g., A100, RTX 3090, RTX 4090, H100)."
+            )
+        },
+    )
+    load_teacher_in_8bit: bool = field(default=False, metadata={"help": "Use 8 bit precision for the teacher model."})
+    load_teacher_in_4bit: bool = field(default=False, metadata={"help": "Use 4 bit precision for the teacher model."})
+    load_student_in_8bit: bool = field(default=False, metadata={"help": "Use 8 bit precision for the student model."})
+    load_student_in_4bit: bool = field(default=False, metadata={"help": "Use 4 bit precision for the student model."})
+    bnb_4bit_quant_type: Optional[str] = field(
+        default="nf4", metadata={"help": "Quantization type if the teacher is quantized (fp4 or nf4)"}
+    )
+    use_bnb_nested_quant: bool = field(default=False, metadata={"help": "Whether or not to use nested quantization."})
+    lora_r: Optional[int] = field(
+        default=16,
+        metadata={"help": "LoRA R value."},
+    )
+    lora_alpha: Optional[int] = field(
+        default=32,
+        metadata={"help": "LoRA alpha."},
+    )
+    lora_dropout: Optional[float] = field(
+        default=0.05,
+        metadata={"help": "LoRA dropout."},
+    )
+    lora_target_modules: Optional[List[str]] = field(
+        default=None,
+        metadata={"help": "LoRA target modules."},
+    )
+    lora_modules_to_save: Optional[List[str]] = field(
+        default=None,
+        metadata={"help": "Model layers to unfreeze & train"},
+    )
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    train_dataset_name: List[str] = field(
+        default=None,
+        metadata={
+            "help": "The name of the training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load LibriSpeech "
+            "and Common Voice, set `train_dataset_name='librispeech_asr+common_voice'`."
+        },
+    )
+    train_dataset_config_name: Optional[List[str]] = field(
+        default=None,
+        metadata={
+            "help": "The configuration name of the training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset configs by a '+' symbol. Note that the order of the configs should "
+            "match the order of the datasets."
+        },
+    )
+    train_dataset_samples: Optional[List[str]] = field(
+        default=None,
+        metadata={
+            "help": "Number of samples in each dataset when loading multiple datasets with streaming mode. "
+            "Not required when using one dataset or non-streaming mode. The sample values provide the sampling "
+            "probability for each dataset. Setting them equal to the number of sample values ensures that every "
+            "sample from every dataset is used once per epoch."
+        },
+    )
+    eval_dataset_name: Optional[List[str]] = field(
+        default=None,
+        metadata={
+            "help": "The name of the evaluation dataset to use (via the datasets library). Defaults to the training "
+            "dataset name if unspecified. Load multiple evaluation datasets by separating dataset "
+            "ids by a '+' symbol."
+        },
+    )
+    eval_dataset_config_name: Optional[List[str]] = field(
+        default=None,
+        metadata={
+            "help": "The configuration name of the evaluation dataset to use (via the datasets library). Defaults to the "
+            "training dataset config name if unspecified."
+        },
+    )
+    dataset_cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to cache directory for saving and loading datasets"},
+    )
+    overwrite_cache: bool = field(
+        default=False,
+        metadata={"help": "Overwrite the cached training and evaluation sets"},
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing if using non-streaming mode."},
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of evaluation examples to this value if set."
+            )
+        },
+    )
+    text_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the generated text data in the training set."},
+    )
+    prompt_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the prompt data. Defaults to 'prompt'"},
+    )
+    eval_text_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the generated text data in the evaluation set."},
+    )
+    eval_prompt_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the prompt data in the evaluation set."},
+    )
+    max_label_length: int = field(
+        default=2048,
+        metadata={"help": "Truncate target labels that are longer `max_label_length` tokens."},
+    )
+    pad_target_to_multiple_of: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "If set will pad the target sequence to a multiple of the provided value. This is important to "
+                "avoid triggering recompilations when using torch compile. If unspecified, will default to padding "
+                "the targets to max length."
+            )
+        },
+    )
+    preprocessing_only: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether to only do data preprocessing and skip training. This is especially useful when data "
+                "preprocessing errors out in distributed training due to timeout. In this case, one should run the "
+                "preprocessing in a non-distributed setup with `preprocessing_only=True` so that the cached datasets "
+                "can consequently be loaded in distributed training"
+            )
+        },
+    )
+    train_split_name: Optional[List[str]] = field(
+        default=lambda: ["train"],
+        metadata={
+            "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
+        },
+    )
+    eval_split_name: Optional[List[str]] = field(
+        default=lambda: ["validation"],
+        metadata={
+            "help": (
+                "The name of the evaluation data set split to use (via the datasets library). Defaults to 'validation'"
+            )
+        },
+    )
+    streaming: bool = field(
+        default=False,
+        metadata={"help": "Whether to use Datasets' streaming mode to load and pre-process the data."},
+    )
+    wandb_project: str = field(
+        default="distil-mixtral",
+        metadata={"help": "The name of the wandb project."},
+    )
+
+
+@dataclass
+class DistillationTrainingArguments(Seq2SeqTrainingArguments):
+    freeze_lm_head: Optional[bool] = field(
+        default=False, metadata={"help": "Whether to freeze the LM head of the student model."}
+    )
+    temperature: Optional[float] = field(
+        default=2.0, metadata={"help": "Temperature to anneal the logits when computing the softmax."}
+    )
+    kl_weight: Optional[float] = field(
+        default=1.0,
+        metadata={
+            "help": (
+                "Weighting assigned to the MSE loss in the KD formulation. MSE loss is "
+                "computed between the teacher-student hidden states and attentions."
+            )
+        },
+    )
+    output_router_logits: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether or not to return the router logits in the forward pass. Enabling this will "
+            "also configure the model to compute the auxiliary loss."
+        },
+    )
+    dtype: Optional[str] = field(
+        default="float32",
+        metadata={
+            "help": (
+                "The data type (dtype) in which to run training. One of `float32` (full-precision), "
+                "`float16` or `bfloat16` (both half-precision)."
+            )
+        },
+    )
+
+
+@dataclass
+class DataCollatorCausalLMWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received.
+    Args:
+        tokenizer ([`PreTrainedTokenizer`])
+            The tokenizer used for tokenizing the data.
+        target_padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned target sequences (according to the model's padding side and padding index).
+            See above for details.
+        max_target_length (:obj:`int`, `optional`):
+            Maximum length of the ``labels`` of the returned list and optionally padding length (see above).
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    target_padding: Union[bool, str] = "max_length"
+    max_target_length: Optional[int] = None
+
+    def __call__(self, features: List[Dict[str, Union[List[int], np.ndarray]]]) -> BatchEncoding:
+        # dataloader returns a list of features which we convert to a dict
+        label_features = {"input_ids": [feature["labels"] for feature in features]}
+        prompt_lengths = [feature["prompt_length"] for feature in features]
+
+        batch = self.tokenizer.pad(
+            label_features,
+            max_length=self.max_target_length,
+            padding=self.target_padding,
+            return_tensors="pt",
+        )
+
+        labels_mask = batch["attention_mask"]
+
+        # don't include prompts in loss calculation
+        for idx in range(len(prompt_lengths)):
+            labels_mask[idx, : prompt_lengths[idx]] = 0
+
+        # replace padding with -100 to ignore loss correctly
+        labels = batch["input_ids"].masked_fill(labels_mask.ne(1), -100)
+
+        batch["labels"] = labels
+
+        return batch
+
+
+def log_metric(
+    accelerator,
+    metrics: Dict,
+    train_time: float,
+    step: int,
+    epoch: int,
+    learning_rate: float = None,
+    prefix: str = "train",
+):
+    """Helper function to log all training/evaluation metrics with the correct prefixes and styling."""
+    log_metrics = {}
+    for k, v in metrics.items():
+        log_metrics[f"{prefix}/{k}"] = v
+    log_metrics[f"{prefix}/time"] = train_time
+    log_metrics[f"{prefix}/epoch"] = epoch
+    if learning_rate is not None:
+        log_metrics[f"{prefix}/learning_rate"] = learning_rate
+    accelerator.log(log_metrics, step=step)
+
+
+def log_pred(
+    accelerator,
+    pred_str: List[str],
+    label_str: List[str],
+    step: int,
+    epoch: int,
+    evaluation_strategy: str,
+    prefix: str = "eval",
+    num_lines: int = 200000,
+):
+    """Helper function to log target/predicted transcriptions to weights and biases (wandb)."""
+    if accelerator.is_main_process:
+        wandb_tracker = accelerator.get_tracker("wandb")
+        # pretty name for current step: step 50000 -> step 50k
+        cur_step_pretty = f"{int(step // 1000)}k" if step > 1000 else step
+        prefix_pretty = prefix.replace("/", "-")
+
+        if evaluation_strategy == "epoch":
+            table_name = f"predictions/{prefix_pretty}-epoch-{epoch}"
+        else:
+            table_name = f"predictions/{prefix_pretty}-step-{cur_step_pretty}"
+
+        # convert str data to a wandb compatible format
+        str_data = [[label_str[i], pred_str[i]] for i in range(len(pred_str))]
+        # log as a table with the appropriate headers
+        wandb_tracker.log_table(
+            table_name=table_name,
+            columns=["Target", "Pred"],
+            data=str_data[:num_lines],
+            step=step,
+        )
+
+
+def convert_dataset_str_to_list(
+    dataset_names,
+    dataset_config_names,
+    splits=None,
+    text_column_names=None,
+    prompt_column_names=None,
+    dataset_samples=None,
+    default_split="train",
+) -> List[Dict]:
+    """
+    Given three lists of dataset names, configs and splits, this function groups the corresponding
+    names/configs/splits. Each dataset is assigned a unique dictionary with these metadata values, and the
+    function returns a list of dictionaries, one for each dataset.
+    """
+    if isinstance(dataset_names, str):
+        dataset_names = [dataset_names]
+        splits = [splits] if splits else None
+        text_column_names = [text_column_names] if text_column_names else None
+        prompt_column_names = [prompt_column_names] if prompt_column_names else None
+    if isinstance(dataset_config_names, str):
+        dataset_config_names = [dataset_config_names]
+
+    if len(dataset_names) == 1 and len(dataset_config_names) > 1:
+        dataset_names = len(dataset_config_names) * dataset_names
+
+    if isinstance(splits, list) and len(splits) == 1 and len(dataset_config_names) > 1:
+        splits = len(dataset_config_names) * splits
+
+    # basic checks to ensure we've got the right number of datasets/configs/splits/columns/probs
+    if dataset_config_names is not None and len(dataset_names) != len(dataset_config_names):
+        raise ValueError(
+            f"Ensure one config is passed for each dataset, got {len(dataset_names)} datasets and"
+            f" {len(dataset_config_names)} configs."
+        )
+
+    if splits is not None and len(splits) != len(dataset_names):
+        raise ValueError(
+            f"Ensure one split is passed for each dataset, got {len(dataset_names)} datasets and {len(splits)} splits."
+        )
+
+    if text_column_names is not None and len(text_column_names) != len(dataset_names):
+        raise ValueError(
+            f"Ensure one text column name is passed for each dataset, got {len(dataset_names)} datasets and"
+            f" {len(text_column_names)} text column names."
+        )
+
+    if prompt_column_names is not None and len(prompt_column_names) != len(dataset_names):
+        raise ValueError(
+            f"Ensure one prompt column name is passed for each dataset, got {len(dataset_names)} datasets and"
+            f" {len(prompt_column_names)} prompt column names."
+        )
+
+    if dataset_samples is not None:
+        if len(dataset_samples) != len(dataset_names):
+            raise ValueError(
+                f"Ensure one sample is passed for each dataset, got {len(dataset_names)} datasets and "
+                f"{len(dataset_samples)} samples."
+            )
+        dataset_samples = [float(ds_sample) for ds_sample in dataset_samples]
+    else:
+        dataset_samples = [None] * len(dataset_names)
+
+    dataset_config_names = (
+        dataset_config_names if dataset_config_names is not None else ["default" for _ in range(len(dataset_names))]
+    )
+    text_column_names = (
+        text_column_names if text_column_names is not None else ["text" for _ in range(len(dataset_names))]
+    )
+    prompt_column_names = (
+        prompt_column_names if prompt_column_names is not None else ["prompt" for _ in range(len(dataset_names))]
+    )
+    splits = splits if splits is not None else [default_split for _ in range(len(dataset_names))]
+
+    dataset_names_dict = []
+    for i, ds_name in enumerate(dataset_names):
+        dataset_names_dict.append(
+            {
+                "name": ds_name,
+                "config": dataset_config_names[i],
+                "split": splits[i],
+                "text_column_name": text_column_names[i],
+                "prompt_column_name": prompt_column_names[i],
+                "samples": dataset_samples[i],
+            }
+        )
+    return dataset_names_dict
+
+
+def load_multiple_datasets(
+    dataset_names: Union[List, str],
+    dataset_config_names: Union[List, str],
+    splits: Optional[Union[List, str]] = None,
+    text_column_names: Optional[List] = None,
+    prompt_column_names: Optional[List] = None,
+    stopping_strategy: Optional[str] = "first_exhausted",
+    dataset_samples: Optional[Union[List, np.array]] = None,
+    streaming: Optional[bool] = False,
+    seed: Optional[int] = None,
+    accelerator: Optional[Accelerator] = None,
+    **kwargs,
+) -> Union[Dataset, IterableDataset]:
+    dataset_names_dict = convert_dataset_str_to_list(
+        dataset_names, dataset_config_names, splits, text_column_names, prompt_column_names, dataset_samples
+    )
+
+    if dataset_samples is not None:
+        dataset_samples = [ds_dict["samples"] for ds_dict in dataset_names_dict]
+        probabilities = np.array(dataset_samples) / np.sum(dataset_samples)
+    else:
+        probabilities = None
+
+    all_datasets = []
+    # iterate over the datasets we want to interleave
+    for dataset_dict in tqdm(
+        dataset_names_dict,
+        desc="Combining datasets...",
+        disable=not accelerator.is_main_process,
+    ):
+        dataset = load_dataset(
+            dataset_dict["name"],
+            dataset_dict["config"],
+            split=dataset_dict["split"],
+            streaming=streaming,
+            **kwargs,
+        )
+
+        columns_to_keep = {"text"}
+        dataset_features = dataset.features.keys()
+
+        if dataset_dict["text_column_name"] not in dataset_features:
+            raise ValueError(
+                f"Text column name {dataset_dict['text_column_name']} not found in dataset"
+                f" '{dataset_dict['name']}'. Make sure to set `--text_column_name` to the"
+                f" correct text column - one of {', '.join(dataset_features)}."
+            )
+
+        # blanket renaming of all transcription columns to text
+        if dataset_dict["text_column_name"] != "text":
+            dataset = dataset.rename_column(dataset_dict["text_column_name"], "text")
+
+        # blanket renaming of all prompt columns to prompt
+        if dataset_dict["prompt_column_name"] is not None:
+            if dataset_dict["prompt_column_name"] not in dataset_features:
+                raise ValueError(
+                    f"Prompt column name {dataset_dict['prompt_column_name']} not found in dataset"
+                    f" '{dataset_dict['name']}'. Make sure to set `--prompt_column_name` to the"
+                    f" correct prompt column - one of {', '.join(dataset_features)}."
+                )
+            elif dataset_dict["prompt_column_name"] != "prompt":
+                dataset = dataset.rename_column(dataset_dict["prompt_column_name"], "prompt")
+            columns_to_keep.add("prompt")
+
+        dataset = dataset.remove_columns(set(dataset_features - columns_to_keep))
+        all_datasets.append(dataset)
+
+    if len(all_datasets) == 1:
+        # we have a single dataset so just return it as is
+        return all_datasets[0]
+
+    if streaming:
+        interleaved_dataset = interleave_datasets(
+            all_datasets,
+            stopping_strategy=stopping_strategy,
+            probabilities=probabilities,
+            seed=seed,
+        )
+    else:
+        interleaved_dataset = concatenate_datasets(all_datasets)
+
+    return interleaved_dataset
+
+
+def sorted_checkpoints(output_dir=None, checkpoint_prefix="checkpoint") -> List[str]:
+    """Helper function to sort saved checkpoints from oldest to newest."""
+    ordering_and_checkpoint_path = []
+
+    glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}-*") if os.path.isdir(x)]
+
+    for path in glob_checkpoints:
+        regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path)
+        if regex_match is not None and regex_match.groups() is not None:
+            ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))
+
+    checkpoints_sorted = sorted(ordering_and_checkpoint_path)
+    checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
+    return checkpoints_sorted
+
+
+def rotate_checkpoints(save_total_limit=None, output_dir=None, checkpoint_prefix="checkpoint") -> None:
+    """Helper function to delete old checkpoints."""
+    if save_total_limit is None or save_total_limit <= 0:
+        return
+    # Check if we should delete older checkpoint(s)
+    checkpoints_sorted = sorted_checkpoints(output_dir=output_dir, checkpoint_prefix=checkpoint_prefix)
+    if len(checkpoints_sorted) <= save_total_limit:
+        return
+
+    number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
+    checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
+    for checkpoint in checkpoints_to_be_deleted:
+        logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
+        shutil.rmtree(checkpoint, ignore_errors=True)
+
+
+_RE_CHECKPOINT = re.compile(r"^checkpoint-(\d+)-epoch-(\d+)$")
+
+
+def get_last_checkpoint(folder):
+    content = os.listdir(folder)
+    checkpoints = [
+        path
+        for path in content
+        if _RE_CHECKPOINT.search(path) is not None and os.path.isdir(os.path.join(folder, path))
+    ]
+    if len(checkpoints) == 0:
+        return
+    return os.path.join(folder, max(checkpoints, key=lambda x: int(_RE_CHECKPOINT.search(x).groups()[0])))
+
+
+def get_parameter_names(model, forbidden_layer_types, forbidden_module=None):
+    """
+    Returns the names of the model parameters that are not inside a forbidden layer or forbidden module.
+    Can be used to get a subset of parameter names for decay masks, or to exclude parameters from an optimiser
+    (e.g. if the module is frozen).
+    """
+    result = []
+    for name, child in model.named_children():
+        result += [
+            f"{name}.{n}"
+            for n in get_parameter_names(child, forbidden_layer_types, forbidden_module)
+            if not (
+                isinstance(child, tuple(forbidden_layer_types))
+                or (child in tuple(forbidden_module) if forbidden_module is not None else False)
+            )
+        ]
+    # Add model specific parameters (defined with nn.Parameter) since they are not in any child.
+    result += list(model._parameters.keys())
+    return result
+
+
+def get_quantization_config(
+    model_args: ModelArguments, torch_dtype: torch.dtype
+) -> tuple[BitsAndBytesConfig | None, BitsAndBytesConfig | None]:
+    if model_args.load_teacher_in_4bit:
+        quantization_config_teacher = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch_dtype,
+            bnb_4bit_quant_type=model_args.bnb_4bit_quant_type,
+            bnb_4bit_use_double_quant=model_args.use_bnb_nested_quant,
+        )
+    elif model_args.load_teacher_in_8bit:
+        quantization_config_teacher = BitsAndBytesConfig(load_in_8bit=True)
+    else:
+        quantization_config_teacher = None
+
+    if model_args.load_student_in_4bit:
+        quantization_config_student = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch_dtype,
+            bnb_4bit_quant_type=model_args.bnb_4bit_quant_type,
+            bnb_4bit_use_double_quant=model_args.use_bnb_nested_quant,
+        )
+    elif model_args.load_student_in_8bit:
+        quantization_config_student = BitsAndBytesConfig(load_in_8bit=True)
+    else:
+        quantization_config_student = None
+
+    return quantization_config_teacher, quantization_config_student
+
+
+def main():
+    # 1. Parse input arguments
+    # We keep distinct sets of args, for cleaner separation of model/data/training related args
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, DistillationTrainingArguments))
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
+        # If we pass only one argument to the script and it's the path to a yaml file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_yaml_file(yaml_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # 2. Initialize the accelerator
+    # We will let the accelerator handle device placement for us in this example
+    # We simply have to specify the training precision and any trackers being used
+    # We'll use the same dtype arguments as our JAX/Flax training script and convert
+    # it to accelerate format
+    if training_args.dtype == "float16":
+        mixed_precision = "fp16"
+        teacher_dtype = torch.float16
+    elif training_args.dtype == "bfloat16":
+        mixed_precision = "bf16"
+        teacher_dtype = torch.bfloat16
+    else:
+        mixed_precision = "no"
+        teacher_dtype = torch.float32
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
+        mixed_precision=mixed_precision,
+        log_with=training_args.report_to,
+        project_dir=training_args.output_dir,
+    )
+
+    accelerator.init_trackers(project_name=data_args.wandb_project)
+
+    # 3. Set-up basic logging
+    # Create one log on every process with the configuration for debugging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    # Log a small summary on each proces
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
+        f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
+    )
+
+    # Set the verbosity to info of the Transformers logger (on main process only)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # 4. Detecting last checkpoint and eventually continue from last checkpoint
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+
+    # 5. Handle the repository creation
+    if accelerator.is_main_process:
+        if training_args.push_to_hub:
+            if training_args.hub_model_id is None:
+                repo_name = get_full_repo_name(
+                    Path(training_args.output_dir).absolute().name,
+                    token=training_args.hub_token,
+                )
+            else:
+                repo_name = training_args.hub_model_id
+            create_repo(repo_name, exist_ok=True, token=training_args.hub_token)
+
+            with open(os.path.join(training_args.output_dir, ".gitignore"), "w+") as gitignore:
+                if "wandb" not in gitignore:
+                    gitignore.write("wandb\n")
+        elif training_args.output_dir is not None:
+            os.makedirs(training_args.output_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    # 6. Load dataset - either streaming or non-streaming (offline)
+    raw_datasets = IterableDatasetDict() if data_args.streaming else DatasetDict()
+
+    # set seed for determinism
+    set_seed(training_args.seed)
+
+    if training_args.do_train:
+        raw_datasets["train"] = load_multiple_datasets(
+            data_args.train_dataset_name,
+            data_args.train_dataset_config_name,
+            splits=data_args.train_split_name,
+            text_column_names=data_args.text_column_name,
+            prompt_column_names=data_args.prompt_column_name,
+            streaming=data_args.streaming,
+            dataset_samples=data_args.train_dataset_samples,
+            seed=training_args.seed,
+            accelerator=accelerator,
+            cache_dir=data_args.dataset_cache_dir,
+            token=model_args.token,
+        )
+        raw_datasets_train_features = set(raw_datasets["train"].features.keys())
+
+    if training_args.do_eval:
+        dataset_names_dict = convert_dataset_str_to_list(
+            data_args.eval_dataset_name if data_args.eval_dataset_name else data_args.train_dataset_name,
+            (
+                data_args.eval_dataset_config_name
+                if data_args.eval_dataset_config_name
+                else data_args.train_dataset_config_name
+            ),
+            splits=data_args.eval_split_name,
+            text_column_names=data_args.eval_text_column_name,
+            prompt_column_names=data_args.eval_prompt_column_name,
+        )
+        all_eval_splits = []
+        if len(dataset_names_dict) == 1:
+            # load a single eval set
+            dataset_dict = dataset_names_dict[0]
+            all_eval_splits.append("eval")
+            raw_datasets["eval"] = load_dataset(
+                dataset_dict["name"],
+                dataset_dict["config"],
+                split=dataset_dict["split"],
+                cache_dir=data_args.dataset_cache_dir,
+                token=model_args.token,
+                streaming=data_args.streaming,
+            )
+            if dataset_dict["text_column_name"] != "text":
+                raw_datasets["eval"] = raw_datasets["eval"].rename_column(data_args.eval_text_column_name, "text")
+            if dataset_dict["prompt_column_name"] != "prompt":
+                raw_datasets["eval"] = raw_datasets["eval"].rename_column(data_args.eval_prompt_column_name, "prompt")
+        else:
+            # load multiple eval sets
+            for dataset_dict in dataset_names_dict:
+                pretty_name = f"{dataset_dict['name'].split('/')[-1]}/{dataset_dict['config'].replace('.', '-')}"
+                all_eval_splits.append(pretty_name)
+                raw_datasets[pretty_name] = load_dataset(
+                    dataset_dict["name"],
+                    dataset_dict["config"],
+                    split=dataset_dict["split"],
+                    cache_dir=data_args.dataset_cache_dir,
+                    token=model_args.token,
+                    streaming=data_args.streaming,
+                )
+                # make column names consistent (text, audio)
+                if dataset_dict["text_column_name"] != "text":
+                    raw_datasets[pretty_name] = raw_datasets[pretty_name].rename_column(
+                        dataset_dict["text_column_name"], "text"
+                    )
+                if dataset_dict["prompt_column_name"] != "prompt":
+                    raw_datasets[pretty_name] = raw_datasets[pretty_name].rename_column(
+                        dataset_dict["prompt_column_name"], "prompt"
+                    )
+                raw_datasets[pretty_name] = raw_datasets[pretty_name].remove_columns(
+                    set(raw_datasets[pretty_name].features.keys()) - {"text", "prompt"}
+                )
+
+    if not training_args.do_train and not training_args.do_eval:
+        raise ValueError(
+            "Cannot not train and not do evaluation. At least one of training or evaluation has to be performed."
+        )
+
+    # 7. Load pretrained model, tokenizer, and feature extractor
+    config = AutoConfig.from_pretrained(
+        (model_args.config_name if model_args.config_name else model_args.model_name_or_path),
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        token=model_args.token,
+    )
+    if training_args.output_router_logits:
+        config.output_router_logits = True
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        (model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path),
+        cache_dir=model_args.cache_dir,
+        use_fast=model_args.use_fast_tokenizer,
+        revision=model_args.model_revision,
+        token=model_args.token,
+    )
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    quantization_config_teacher, quantization_config_student = get_quantization_config(
+        model_args, torch_dtype=teacher_dtype
+    )
+
+    # The teacher model can safely be cast to the dtype of training since we don't
+    # update the params
+    teacher_model = AutoModelForCausalLM.from_pretrained(
+        model_args.teacher_model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=model_args.token,
+        low_cpu_mem_usage=True,
+        torch_dtype=teacher_dtype,
+        attn_implementation=model_args.attn_implementation,
+        quantization_config=quantization_config_teacher,
+    )
+
+    student_model = AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path,
+        config=config,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        subfolder=model_args.subfolder,
+        token=model_args.token,
+        torch_dtype=teacher_dtype,
+        low_cpu_mem_usage=True,
+        attn_implementation=model_args.attn_implementation,
+        quantization_config=quantization_config_student,
+    )
+
+    if quantization_config_student is not None:
+        lora_config = LoraConfig(
+            r=model_args.lora_r,
+            lora_alpha=model_args.lora_alpha,
+            target_modules=model_args.lora_target_modules,
+            lora_dropout=model_args.lora_dropout,
+            bias="none",
+            task_type="CAUSAL_LM",
+        )
+        student_model = get_peft_model(student_model, lora_config)
+
+    if student_model.generation_config.bos_token_id is None or teacher_model.generation_config.bos_token_id is None:
+        raise ValueError(
+            f"Make sure that `generation_config.bos_token_id` is correctly defined for both the "
+            f"student and teacher model. Got {student_model.generation_config.bos_token_id} for the "
+            f"student and {teacher_model.generation_config.bos_token_id} for the teacher."
+        )
+
+    # enable gradient checkpointing if necessary
+    if training_args.gradient_checkpointing:
+        student_model.gradient_checkpointing_enable()
+
+    def set_trainable_parameters(module, requires_grad=False):
+        for param in module.parameters():
+            param.requires_grad = requires_grad
+        module._requires_grad = requires_grad
+
+    # freeze student lm head if necessary
+    if training_args.freeze_lm_head:
+        set_trainable_parameters(student_model.lm_head, requires_grad=False)
+        # TODO(SG): possibly upgrade this to an error
+        if training_args.gradient_checkpointing:
+            logger.warning(
+                "Freezing the LM head is not compatible with gradient checkpointing. Set `--gradient_checkpointing=False`, "
+                "or un-freeze the LM head with `--freeze_lm_head=False`. Overriding gradient checkpointing to False."
+            )
+
+    student_model.generation_config.max_length = data_args.max_label_length
+
+    # 8. Save all pre-processed tokenizers/config/generation configs
+    if accelerator.is_main_process:
+        tokenizer.save_pretrained(training_args.output_dir)
+        # save the config and generation config as well
+        config.save_pretrained(training_args.output_dir)
+        student_model.generation_config.save_pretrained(training_args.output_dir)
+
+    accelerator.wait_for_everyone()
+
+
+    # 10. Preprocessing the datasets: we need to combine the prompt and generations and tokenize the targets.
+    # 10.1: Define the pre-processing constants
+    max_label_length = (
+        data_args.max_label_length if data_args.max_label_length is not None else config.max_length
+    )
+    num_workers = data_args.preprocessing_num_workers
+    dataloader_num_workers = training_args.dataloader_num_workers
+    prefetch_factor = training_args.dataloader_prefetch_factor
+    eos_token_id = tokenizer.eos_token_id
+
+    # 10.2: filter based on maximum number of training/evaluation samples
+    if training_args.do_train and data_args.max_train_samples is not None:
+        raw_datasets["train"] = (
+            raw_datasets["train"].take(data_args.max_train_samples)
+            if data_args.streaming
+            else raw_datasets["train"].select(range(data_args.max_train_samples))
+        )
+
+    if training_args.do_eval and data_args.max_eval_samples is not None:
+        for eval_split in all_eval_splits:
+            raw_datasets[eval_split] = (
+                raw_datasets[eval_split].take(data_args.max_eval_samples)
+                if data_args.streaming
+                else raw_datasets[eval_split].select(range(data_args.max_eval_samples))
+            )
+
+    # 10.3: pre-process training/evaluation datasets
+    def prepare_datasets(example):
+        prompt_ids = tokenizer(example["prompt"]).input_ids
+        gen_ids = tokenizer(example["text"], add_special_tokens=False).input_ids + [eos_token_id]
+        if prompt_ids[-1] == eos_token_id:
+            prompt_ids = prompt_ids[:-1]
+        example["labels"] = prompt_ids + gen_ids
+        example["prompt_length"] = len(prompt_ids)
+        return example
+
+    vectorized_datasets = IterableDatasetDict() if data_args.streaming else DatasetDict()
+    if training_args.do_train:
+        # with streaming mode we can only have 1 worker, whereas with non-streaming
+        # we can use `num_workers` (which is much faster)
+        # We gate the pre-processing function accordingly
+        map_fn_train = partial(
+            raw_datasets["train"].map,
+            function=prepare_datasets,
+            remove_columns=raw_datasets_train_features,
+        )
+        with accelerator.main_process_first():
+            vectorized_datasets["train"] = (
+                map_fn_train(num_proc=num_workers, desc="preprocess train dataset")
+                if not data_args.streaming
+                else map_fn_train()
+            )
+    if training_args.do_eval:
+        for eval_split in all_eval_splits:
+            raw_datasets_eval_features = list(raw_datasets[eval_split].features.keys())
+            map_fn_eval = partial(
+                raw_datasets[eval_split].map, function=prepare_datasets, remove_columns=raw_datasets_eval_features
+            )
+            with accelerator.main_process_first():
+                vectorized_datasets[eval_split] = (
+                    map_fn_eval(num_proc=num_workers, desc="preprocess eval dataset")
+                    if not data_args.streaming
+                    else map_fn_eval()
+                )
+
+    # 10.4: Filter training data with labels longer than `max_label_length`
+    def is_labels_in_length_range(labels):
+        return 0 < len(labels) <= max_label_length
+
+    filter_by_labels_fn = partial(
+        vectorized_datasets.filter, function=is_labels_in_length_range, input_columns=["labels"]
+    )
+    with accelerator.main_process_first():
+        vectorized_datasets = (
+            filter_by_labels_fn(num_proc=num_workers, desc="filtering train dataset")
+            if not data_args.streaming
+            else filter_by_labels_fn()
+        )
+
+    # Pre-processing complete!
+    # For large datasets it is advised to run the preprocessing on a
+    # single machine first with `--preprocessing_only` since there will mostly likely
+    # be a timeout when running the script in distributed mode.
+    # In a second step, `--preprocessing_only` can then be set to `False` to load the
+    # cached dataset
+    if data_args.preprocessing_only:
+        if data_args.streaming:
+            raise ValueError(
+                "When using streaming mode, dataset pre-processing is performed on the fly, hence there is no notion"
+                "of a cached pre-processed dataset. Remove the argument `--preprocessing_only` to run pre-processing "
+                "on the fly with streaming mode."
+            )
+        cache = {k: v.cache_files for k, v in vectorized_datasets.items()}
+        logger.info(f"Data preprocessing finished. Files cached at {cache}.")
+        return
+
+    # 11. Define Evaluation Metrics
+    def compute_metrics(preds, labels):
+        # TODO(SG): better metrics for performance?
+        # replace padded labels by the padding token
+        for idx in range(len(labels)):
+            labels[idx][labels[idx] == -100] = tokenizer.pad_token_id
+        pred_str = tokenizer.batch_decode(preds, skip_special_tokens=True)
+        label_str = tokenizer.batch_decode(labels, skip_special_tokens=True)
+        return pred_str, label_str
+
+    # 12. Define Training Schedule
+    # 12.1: Store some constants
+    per_device_train_batch_size = int(training_args.per_device_train_batch_size)
+    train_batch_size = per_device_train_batch_size * accelerator.num_processes
+    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
+    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
+
+    # 12.2: Set max training steps
+    if not data_args.streaming and training_args.max_steps < 0:
+        num_epochs = int(training_args.num_train_epochs)
+        steps_per_epoch = len(vectorized_datasets["train"]) // (train_batch_size * gradient_accumulation_steps)
+        total_train_steps = steps_per_epoch * num_epochs
+    elif training_args.max_steps > 0:
+        logger.info("max_steps is given, it will override any value given in num_train_epochs")
+        total_train_steps = int(training_args.max_steps)
+        if not data_args.streaming:
+            steps_per_epoch = len(vectorized_datasets["train"]) // (train_batch_size * gradient_accumulation_steps)
+            num_epochs = int(np.ceil(total_train_steps / steps_per_epoch))
+        else:
+            # Setting a very large number of epochs so we go as many times as necessary over the iterator.
+            num_epochs = sys.maxsize
+            steps_per_epoch = total_train_steps
+    else:
+        raise ValueError("max_steps must be specified when training with a streaming (iterable) dataset")
+
+    # 12.3: Set evaluation steps
+    if training_args.evaluation_strategy == "epoch":
+        eval_steps = steps_per_epoch
+    elif training_args.eval_steps is None:
+        logger.info(
+            f"eval_steps is not set, evaluating at the end of {'each epoch' if not data_args.streaming else 'training'}"
+        )
+        eval_steps = steps_per_epoch
+    else:
+        eval_steps = training_args.eval_steps
+
+    # 12.4: Set save steps
+    if training_args.save_strategy == "epoch":
+        save_steps = steps_per_epoch
+    elif training_args.save_strategy == "steps":
+        save_steps = training_args.save_steps
+    else:
+        save_steps = sys.maxsize
+
+    # 13. Define optimizer, LR scheduler, collator
+    decay_parameters = get_parameter_names(
+        student_model,
+        [nn.LayerNorm],
+    )
+    decay_parameters = [name for name in decay_parameters if "bias" not in name]
+    optimizer_grouped_parameters = [
+        {
+            "params": [param for name, param in student_model.named_parameters() if name in decay_parameters],
+            "weight_decay": training_args.weight_decay,
+        },
+        {
+            "params": [param for name, param in student_model.named_parameters() if name not in decay_parameters],
+            "weight_decay": 0.0,
+        },
+    ]
+    optimizer = torch.optim.AdamW(
+        params=optimizer_grouped_parameters,
+        lr=training_args.learning_rate,
+        betas=(training_args.adam_beta1, training_args.adam_beta2),
+        eps=training_args.adam_epsilon,
+    )
+
+    # LR scheduler gets stepped by `num_processes` each time -> account for this in warmup / total steps
+    lr_scheduler = get_scheduler(
+        name=training_args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=training_args.warmup_steps * accelerator.num_processes,
+        num_training_steps=total_train_steps * accelerator.num_processes,
+    )
+
+    data_collator = DataCollatorCausalLMWithPadding(
+        tokenizer=tokenizer,
+        target_padding="max_length",
+        max_target_length=max_label_length,
+    )
+
+    # 14. Define generation arguments - we need to do this before we wrap the models in DDP
+    # so that we can still access the configs
+    num_beams = (
+        training_args.generation_num_beams
+        if training_args.generation_num_beams is not None
+        else getattr(student_model.generation_config, "num_beams", 1)
+    )
+
+    # 15. Prepare everything with accelerate
+    student_model, teacher_model, optimizer, lr_scheduler = accelerator.prepare(
+        student_model, teacher_model, optimizer, lr_scheduler
+    )
+
+    def kl_divergence(target_distribution, log_predicted_distribution, labels):
+        kl_loss = nn.KLDivLoss(reduction="none")
+        divergence = kl_loss(log_predicted_distribution, target_distribution)
+        # ignore padded tokens from divergence, i.e. where labels are not set to -100
+        padding_mask = labels >= 0
+        padding_mask = padding_mask.unsqueeze(-1)
+        divergence = divergence * padding_mask
+        # take the average over the mini-batch
+        divergence = divergence.sum() / padding_mask.sum()
+        return divergence
+
+    # Define gradient update step fn
+    def train_step(
+        batch,
+        temperature=2.0,
+    ):
+        student_model.train()
+        teacher_model.eval()
+
+        student_outputs = student_model(**batch)
+        with torch.no_grad():
+            teacher_outputs = teacher_model(**batch)
+
+        # CE (data) loss
+        ce_loss = student_outputs.loss
+        # rescale distribution by temperature to ensure gradients scale correctly
+        teacher_distribution = nn.functional.softmax(teacher_outputs.logits / temperature, dim=-1)
+        # log softmax of student predictions for numerical stability
+        student_distribution = nn.functional.log_softmax(student_outputs.logits / temperature, dim=-1)
+        # KL-divergence loss (scaled by temperature)
+        kl_loss = kl_divergence(teacher_distribution, student_distribution, batch["labels"]) * temperature**2
+
+        # use Distil-Whisper formulation (fix weight of CE loss and tune KL weight)
+        loss = 0.8 * ce_loss + training_args.kl_weight * kl_loss
+        metrics = {"loss": loss, "ce_loss": ce_loss, "kl_loss": kl_loss}
+        return loss, metrics
+
+    # Define eval fn
+    def eval_step(batch):
+        student_model.eval()
+        teacher_model.eval()
+
+        with torch.no_grad():
+            student_outputs = student_model(**batch)
+            teacher_outputs = teacher_model(**batch)
+
+        # CE (data) loss
+        ce_loss = student_outputs.loss
+
+        # log softmax / softmax for numerical stability
+        student_distribution = nn.functional.log_softmax(student_outputs.logits, dim=-1)
+        teacher_distribution = nn.functional.softmax(teacher_outputs.logits, dim=-1)
+        # temperature is always 1 for eval
+        kl_loss = kl_divergence(teacher_distribution, student_distribution, batch["labels"])
+
+        # use Distil-Whisper formulation (fix weight of CE loss and tune KL weight)
+        loss = 0.8 * ce_loss + training_args.kl_weight * kl_loss
+        metrics = {"loss": loss, "ce_loss": ce_loss, "kl_loss": kl_loss}
+        return metrics
+
+    def generate_step(batch):
+        student_model.eval()
+        output_ids = accelerator.unwrap_model(student_model).generate(
+            **batch, max_length=max_label_length, num_beams=num_beams
+        )
+        output_ids = accelerator.pad_across_processes(output_ids, dim=1, pad_index=tokenizer.pad_token_id)
+        return output_ids
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {total_train_steps * train_batch_size * gradient_accumulation_steps}")
+    if not data_args.streaming:
+        logger.info(f"  Num epochs = {num_epochs}")
+    logger.info("  Instantaneous batch size per device =" f" {training_args.per_device_train_batch_size}")
+    logger.info("  Gradient accumulation steps =" f" {gradient_accumulation_steps}")
+    logger.info(
+        f"  Total train batch size (w. parallel & distributed) = {train_batch_size * gradient_accumulation_steps}"
+    )
+    logger.info(f"  Total optimization steps = {total_train_steps}")
+
+    # ======================== Training ================================
+    train_time = 0
+    train_start = time.time()
+    steps_trained_progress_bar = tqdm(
+        range(total_train_steps), desc="Train steps ... ", position=0, disable=not accelerator.is_local_main_process
+    )
+    continue_training = True
+    epochs_trained = 0
+    cur_step = 0
+
+    checkpoint = None
+    if training_args.resume_from_checkpoint is not None:
+        checkpoint = training_args.resume_from_checkpoint
+    elif last_checkpoint is not None:
+        checkpoint = last_checkpoint
+
+    if checkpoint is not None:
+        accelerator.load_state(checkpoint)
+        # Find num steps and epoch from saved state string pattern
+        pattern = r"checkpoint-(\d+)-epoch-(\d+)"
+        match = re.search(pattern, checkpoint)
+        cur_step = int(match.group(1))
+        epochs_trained = int(match.group(2))
+
+        logger.info("  Continuing training from checkpoint, will skip to saved global_step")
+        logger.info(f"  Continuing training from epoch {epochs_trained}")
+        logger.info(f"  Continuing training from global step {cur_step}")
+
+        steps_trained_progress_bar.update(cur_step)
+
+        for epoch in range(0, epochs_trained):
+            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+
+        if not data_args.streaming and training_args.max_steps < 0:
+            # we know exactly the number of steps per epoch, so can skip through the required number of batches
+            resume_step = (cur_step - epochs_trained * steps_per_epoch) * gradient_accumulation_steps
+        else:
+            # Currently we don't know how many steps we've taken in the current epoch
+            # So we just shuffle the dataset one extra time and start from a fresh epoch
+            # This is "good enough" for our purposes but not fully correct
+            resume_step = None
+            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+    else:
+        resume_step = None
+
+    for epoch in range(epochs_trained, num_epochs):
+        vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+        train_dataloader = DataLoader(
+            vectorized_datasets["train"],
+            collate_fn=data_collator,
+            batch_size=per_device_train_batch_size,
+            num_workers=dataloader_num_workers,
+            prefetch_factor=prefetch_factor,
+            pin_memory=training_args.dataloader_pin_memory,
+        )
+        train_dataloader = accelerator.prepare(train_dataloader)
+        if hasattr(train_dataloader, "dataset") and isinstance(train_dataloader.dataset, IterableDataset):
+            train_dataloader.dataset.set_epoch(epoch)
+
+        if resume_step is not None:
+            # Skip the first N batches in the dataloader when resuming from a checkpoint
+            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
+            resume_step = None
+
+        for batch in train_dataloader:
+            with accelerator.accumulate(student_model):
+                loss, train_metric = train_step(batch, temperature=training_args.temperature)
+                accelerator.backward(loss)
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(student_model.parameters(), training_args.max_grad_norm)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            # Check if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                steps_trained_progress_bar.update(1)
+                cur_step += 1
+
+                if cur_step % training_args.logging_steps == 0:
+                    steps_trained_progress_bar.write(
+                        f"Step... ({cur_step} / {total_train_steps} | Loss:"
+                        f" {train_metric['loss']}, Learning Rate:"
+                        f" {lr_scheduler.get_last_lr()[0]})"
+                    )
+                    log_metric(
+                        accelerator,
+                        metrics=train_metric,
+                        learning_rate=lr_scheduler.get_last_lr()[0],
+                        train_time=train_time + time.time() - train_start,
+                        step=cur_step,
+                        epoch=epoch,
+                        prefix="train",
+                    )
+
+                # save checkpoint and weights after each save_steps and at the end of training
+                if (cur_step % save_steps == 0) or cur_step == total_train_steps:
+                    intermediate_dir = os.path.join(training_args.output_dir, f"checkpoint-{cur_step}-epoch-{epoch}")
+                    accelerator.save_state(output_dir=intermediate_dir)
+                    accelerator.wait_for_everyone()
+                    if accelerator.is_main_process:
+                        rotate_checkpoints(training_args.save_total_limit, output_dir=training_args.output_dir)
+                        if training_args.push_to_hub:
+                            upload_folder(
+                                folder_path=training_args.output_dir,
+                                repo_id=repo_name,
+                                repo_type="model",
+                                commit_message=f"Saving train state of step {cur_step}",
+                            )
+
+                if training_args.do_eval and (cur_step % eval_steps == 0 or cur_step == total_train_steps):
+                    train_time += time.time() - train_start
+                    student_model.eval()
+                    # ======================== Evaluating ==============================
+                    for eval_split in all_eval_splits:
+                        eval_metrics = []
+                        eval_preds = []
+                        eval_labels = []
+                        eval_start = time.time()
+
+                        validation_dataloader = DataLoader(
+                            vectorized_datasets[eval_split],
+                            collate_fn=data_collator,
+                            batch_size=per_device_eval_batch_size,
+                            drop_last=False,
+                            num_workers=dataloader_num_workers,
+                            prefetch_factor=prefetch_factor,
+                            pin_memory=training_args.dataloader_pin_memory,
+                        )
+                        validation_dataloader = accelerator.prepare(validation_dataloader)
+
+                        for batch in tqdm(
+                            validation_dataloader,
+                            desc=f"Evaluating {eval_split}...",
+                            position=2,
+                            disable=not accelerator.is_local_main_process,
+                        ):
+                            # Model forward
+                            eval_metric = eval_step(batch)
+                            eval_metric = accelerator.gather_for_metrics(eval_metric)
+                            eval_metrics.append(eval_metric)
+
+                            # generation
+                            if training_args.predict_with_generate:
+                                generated_ids = generate_step(batch)
+                                # Gather all predictions and targets
+                                generated_ids, labels = accelerator.gather_for_metrics(
+                                    (generated_ids, batch["labels"])
+                                )
+                                eval_preds.extend(generated_ids)
+                                eval_labels.extend(labels)
+
+                        eval_time = time.time() - eval_start
+                        # normalize eval metrics
+                        eval_metrics = {
+                            key: torch.mean(torch.stack([d[key] for d in eval_metrics])) for key in eval_metrics[0]
+                        }
+                        try:
+                            eval_metrics["perplexity"] = math.exp(eval_metrics["ce_loss"])
+                        except OverflowError:
+                            eval_metrics["perplexity"] = float("inf")
+
+                        if training_args.predict_with_generate:
+                            pred_str, label_str = compute_metrics(eval_preds, eval_labels)
+                            log_pred(
+                                accelerator,
+                                pred_str,
+                                label_str,
+                                step=cur_step,
+                                epoch=epochs_trained,
+                                evaluation_strategy=training_args.evaluation_strategy,
+                                prefix=eval_split,
+                            )
+
+                        # Print metrics and update progress bar
+                        logger_desc = " ".join([f"Eval {key}: {value} |" for key, value in eval_metrics.items()])
+                        steps_trained_progress_bar.write(
+                            f"Eval results for step ({cur_step} / {total_train_steps} | {logger_desc}"
+                        )
+
+                        log_metric(
+                            accelerator,
+                            metrics=eval_metrics,
+                            train_time=eval_time,
+                            step=cur_step,
+                            epoch=epoch,
+                            prefix=eval_split,
+                        )
+
+                    # flush the train metrics
+                    train_start = time.time()
+
+                # break condition
+                if cur_step == total_train_steps:
+                    # un-wrap student model for save
+                    student_model = accelerator.unwrap_model(student_model)
+                    student_model.save_pretrained(training_args.output_dir)
+                    if training_args.push_to_hub:
+                        upload_folder(
+                            folder_path=training_args.output_dir,
+                            repo_id=repo_name,
+                            repo_type="model",
+                            commit_message=f"Saving final weights of step {cur_step}",
+                        )
+                    continue_training = False
+                    break
+
+        if not continue_training:
+            break
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main()

slurm_job.slurm

@@ -0,0 +1,74 @@
+#!/bin/bash
+#SBATCH --job-name=distil-mixtral
+#SBATCH --nodes=1
+# set 24h for job wall time limit
+#SBATCH --time=4-0:00:00
+#SBATCH --ntasks-per-node=1          # crucial - only 1 task per dist per node!
+#SBATCH --cpus-per-task=32
+#SBATCH --gres=gpu:8
+#SBATCH --exclusive
+#SBATCH --partition=hopper-prod
+#SBATCH --output=/fsx/sanchit/logs/%x-%j.out
+
+set -x -e
+
+source ~/.bashrc
+conda activate venv
+
+echo "START TIME: $(date)"
+
+
+LOG_PATH="/fsx/sanchit/logs/main_log.txt"
+SAVE_DIR="/fsx/sanchit"
+
+GPUS_PER_NODE=8
+NNODES=$SLURM_NNODES
+
+# so processes know who to talk to
+MASTER_ADDR=`scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1`
+
+# From https://i.hsfzxjy.site/2021-03-10-obtain-a-random-unused-tcp-port-with-bash/
+function unused_port() {
+    N=${1:-1}
+    comm -23 \
+        <(seq "1025" "65535" | sort) \
+        <(ss -Htan |
+            awk '{print $4}' |
+            cut -d':' -f2 |
+            sort -u) |
+        shuf |
+        head -n "$N"
+}
+MASTER_PORT=$(unused_port)
+
+# export TORCH_CPP_LOG_LEVEL=INFO
+# export TORCH_DISTRIBUTED_DEBUG=DETAIL
+
+export LAUNCHER="python -u -m accelerate.commands.launch --config_file ./accelerate_config.yaml"
+
+export PROGRAM="./run_distillation.py config_mistral.yaml"
+export CMD="$LAUNCHER $PROGRAM"
+echo $CMD
+
+SRUN_ARGS=" \
+    --wait=60 \
+    --kill-on-bad-exit=1 \
+    "
+
+# py-spy top -s -i -n -- $LAUNCHER --node_rank $SLURM_PROCID --role $SLURMD_NODENAME: $CMD
+clear; srun $SRUN_ARGS --jobid $SLURM_JOB_ID bash -c "$CMD" 2>&1 | tee -a $SAVE_DIR/logs/main_log.txt
+
+
+# srun error handling:
+# --wait=60: wait 60 sec after the first task terminates before terminating all remaining tasks
+# --kill-on-bad-exit=1: terminate a step if any task exits with a non-zero exit code
+
+# SRUN_ARGS=" \
+#     --wait=60 \
+#     --kill-on-bad-exit=1 \
+#     "
+#
+# # py-spy top -s -i -n -- $LAUNCHER --node_rank $SLURM_PROCID --role $SLURMD_NODENAME: $CMD
+# clear; srun $SRUN_ARGS --jobid $SLURM_JOBID bash -c "$CMD" 2>&1 | tee -a $SAVE_DIR/logs/main_log.txt
+
+echo "END TIME: $(date)"

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.model

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

tokenizer_config.json

@@ -0,0 +1,42 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "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
+    }
+  },
+  "additional_special_tokens": [],
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "legacy": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "</s>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

wandb/debug-internal.log

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

wandb/debug.log

@@ -0,0 +1,34 @@
+2024-04-16 20:53:09,244 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Current SDK version is 0.16.1
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Configure stats pid to 965482
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Loading settings from /admin/home/sanchit/.config/wandb/settings
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Loading settings from /fsx/sanchit/mistral-debug-4bit/wandb/settings
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Loading settings from environment variables: {}
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Applying setup settings: {'_disable_service': False}
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Inferring run settings from compute environment: {'program_relpath': 'run_distillation.py', 'program_abspath': '/fsx/sanchit/mistral-debug-4bit/run_distillation.py', 'program': '/fsx/sanchit/mistral-debug-4bit/run_distillation.py'}
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:_log_setup():524] Logging user logs to /fsx/sanchit/mistral-debug-4bit/wandb/run-20240416_205309-xdytsc71/logs/debug.log
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:_log_setup():525] Logging internal logs to /fsx/sanchit/mistral-debug-4bit/wandb/run-20240416_205309-xdytsc71/logs/debug-internal.log
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:init():564] calling init triggers
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:init():571] wandb.init called with sweep_config: {}
+config: {}
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:init():614] starting backend
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:init():618] setting up manager
+2024-04-16 20:53:09,248 INFO    MainThread:965482 [backend.py:_multiprocessing_setup():105] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-04-16 20:53:09,254 INFO    MainThread:965482 [wandb_init.py:init():624] backend started and connected
+2024-04-16 20:53:09,257 INFO    MainThread:965482 [wandb_init.py:init():716] updated telemetry
+2024-04-16 20:53:09,257 INFO    MainThread:965482 [wandb_init.py:init():749] communicating run to backend with 90.0 second timeout
+2024-04-16 20:53:09,461 INFO    MainThread:965482 [wandb_run.py:_on_init():2254] communicating current version
+2024-04-16 20:53:09,520 INFO    MainThread:965482 [wandb_run.py:_on_init():2263] got version response upgrade_message: "wandb version 0.16.6 is available!  To upgrade, please run:\n $ pip install wandb --upgrade"
+
+2024-04-16 20:53:09,520 INFO    MainThread:965482 [wandb_init.py:init():800] starting run threads in backend
+2024-04-16 20:53:15,732 INFO    MainThread:965482 [wandb_run.py:_console_start():2233] atexit reg
+2024-04-16 20:53:15,732 INFO    MainThread:965482 [wandb_run.py:_redirect():2088] redirect: wrap_raw
+2024-04-16 20:53:15,732 INFO    MainThread:965482 [wandb_run.py:_redirect():2153] Wrapping output streams.
+2024-04-16 20:53:15,732 INFO    MainThread:965482 [wandb_run.py:_redirect():2178] Redirects installed.
+2024-04-16 20:53:15,733 INFO    MainThread:965482 [wandb_init.py:init():841] run started, returning control to user process
+2024-04-17 15:43:17,904 INFO    MainThread:965482 [wandb_run.py:_finish():1962] finishing run sanchit-gandhi/distil-mixtral/xdytsc71
+2024-04-17 15:43:17,906 INFO    MainThread:965482 [wandb_run.py:_atexit_cleanup():2202] got exitcode: 0
+2024-04-17 15:43:17,906 INFO    MainThread:965482 [wandb_run.py:_restore():2185] restore
+2024-04-17 15:43:17,907 INFO    MainThread:965482 [wandb_run.py:_restore():2191] restore done
+2024-04-17 15:43:24,105 INFO    MainThread:965482 [wandb_run.py:_footer_history_summary_info():3837] rendering history
+2024-04-17 15:43:24,106 INFO    MainThread:965482 [wandb_run.py:_footer_history_summary_info():3869] rendering summary
+2024-04-17 15:43:24,118 INFO    MainThread:965482 [wandb_run.py:_footer_sync_info():3796] logging synced files

wandb/run-20240416_172306-uygw9yfk/files/conda-environment.yaml

@@ -0,0 +1,300 @@
+name: venv
+channels:
+  - pytorch
+  - nvidia
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - blas=1.0=mkl
+  - brotli-python=1.0.9=py311h6a678d5_7
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2023.12.12=h06a4308_0
+  - certifi=2023.11.17=py311h06a4308_0
+  - cffi=1.16.0=py311h5eee18b_0
+  - cryptography=41.0.7=py311hdda0065_0
+  - cuda-cudart=12.1.105=0
+  - cuda-cupti=12.1.105=0
+  - cuda-libraries=12.1.0=0
+  - cuda-nvrtc=12.1.105=0
+  - cuda-nvtx=12.1.105=0
+  - cuda-opencl=12.3.101=0
+  - cuda-runtime=12.1.0=0
+  - ffmpeg=4.3=hf484d3e_0
+  - filelock=3.13.1=py311h06a4308_0
+  - freetype=2.12.1=h4a9f257_0
+  - giflib=5.2.1=h5eee18b_3
+  - gmp=6.2.1=h295c915_3
+  - gmpy2=2.1.2=py311hc9b5ff0_0
+  - gnutls=3.6.15=he1e5248_0
+  - intel-openmp=2023.1.0=hdb19cb5_46306
+  - jinja2=3.1.2=py311h06a4308_0
+  - jpeg=9e=h5eee18b_1
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - lerc=3.0=h295c915_0
+  - libcublas=12.1.0.26=0
+  - libcufft=11.0.2.4=0
+  - libcufile=1.8.1.2=0
+  - libcurand=10.3.4.101=0
+  - libcusolver=11.4.4.55=0
+  - libcusparse=12.0.2.55=0
+  - libdeflate=1.17=h5eee18b_1
+  - libffi=3.4.4=h6a678d5_0
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgomp=11.2.0=h1234567_1
+  - libiconv=1.16=h7f8727e_2
+  - libidn2=2.3.4=h5eee18b_0
+  - libjpeg-turbo=2.0.0=h9bf148f_0
+  - libnpp=12.0.2.50=0
+  - libnvjitlink=12.1.105=0
+  - libnvjpeg=12.1.1.14=0
+  - libpng=1.6.39=h5eee18b_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libtasn1=4.19.0=h5eee18b_0
+  - libtiff=4.5.1=h6a678d5_0
+  - libunistring=0.9.10=h27cfd23_0
+  - libuuid=1.41.5=h5eee18b_0
+  - libwebp=1.3.2=h11a3e52_0
+  - libwebp-base=1.3.2=h5eee18b_0
+  - llvm-openmp=14.0.6=h9e868ea_0
+  - lz4-c=1.9.4=h6a678d5_0
+  - markupsafe=2.1.1=py311h5eee18b_0
+  - mkl=2023.1.0=h213fc3f_46344
+  - mkl-service=2.4.0=py311h5eee18b_1
+  - mkl_fft=1.3.8=py311h5eee18b_0
+  - mkl_random=1.2.4=py311hdb19cb5_0
+  - mpc=1.1.0=h10f8cd9_1
+  - mpfr=4.0.2=hb69a4c5_1
+  - mpmath=1.3.0=py311h06a4308_0
+  - ncurses=6.4=h6a678d5_0
+  - nettle=3.7.3=hbbd107a_1
+  - networkx=3.1=py311h06a4308_0
+  - numpy=1.26.2=py311h08b1b3b_0
+  - numpy-base=1.26.2=py311hf175353_0
+  - openh264=2.1.1=h4ff587b_0
+  - openjpeg=2.4.0=h3ad879b_0
+  - openssl=3.0.12=h7f8727e_0
+  - pycparser=2.21=pyhd3eb1b0_0
+  - pyopenssl=23.2.0=py311h06a4308_0
+  - pysocks=1.7.1=py311h06a4308_0
+  - python=3.11.5=h955ad1f_0
+  - pytorch-cuda=12.1=ha16c6d3_5
+  - pytorch-mutex=1.0=cuda
+  - pyyaml=6.0.1=py311h5eee18b_0
+  - readline=8.2=h5eee18b_0
+  - requests=2.31.0=py311h06a4308_0
+  - setuptools=68.2.2=py311h06a4308_0
+  - sqlite=3.41.2=h5eee18b_0
+  - sympy=1.12=py311h06a4308_0
+  - tbb=2021.8.0=hdb19cb5_0
+  - tk=8.6.12=h1ccaba5_0
+  - wheel=0.41.2=py311h06a4308_0
+  - xz=5.4.5=h5eee18b_0
+  - yaml=0.2.5=h7b6447c_0
+  - zlib=1.2.13=h5eee18b_0
+  - zstd=1.5.5=hc292b87_0
+  - pip:
+      - absl-py==2.0.0
+      - accelerate==0.29.2
+      - aiohttp==3.9.1
+      - aiosignal==1.3.1
+      - annotated-types==0.6.0
+      - anyio==4.2.0
+      - appdirs==1.4.4
+      - argon2-cffi==23.1.0
+      - argon2-cffi-bindings==21.2.0
+      - arrow==1.3.0
+      - asttokens==2.4.1
+      - astunparse==1.6.3
+      - async-lru==2.0.4
+      - attrs==23.1.0
+      - audioread==3.0.1
+      - babel==2.14.0
+      - beautifulsoup4==4.12.3
+      - bitsandbytes==0.43.1
+      - bleach==6.1.0
+      - cachetools==5.3.2
+      - chardet==5.2.0
+      - charset-normalizer==3.3.2
+      - click==8.1.7
+      - comm==0.2.1
+      - datasets==2.18.1.dev0
+      - debugpy==1.8.1
+      - decorator==5.1.1
+      - deepspeed==0.12.2
+      - defusedxml==0.7.1
+      - dill==0.3.7
+      - docker-pycreds==0.4.0
+      - docstring-parser==0.15
+      - einops==0.7.0
+      - evaluate==0.4.0
+      - executing==2.0.1
+      - fastjsonschema==2.19.1
+      - flatbuffers==23.5.26
+      - fqdn==1.5.1
+      - frozenlist==1.4.1
+      - fsspec==2023.10.0
+      - gast==0.5.4
+      - gitdb==4.0.11
+      - gitpython==3.1.40
+      - google-auth==2.26.1
+      - google-auth-oauthlib==1.2.0
+      - google-pasta==0.2.0
+      - grpcio==1.60.0
+      - h11==0.14.0
+      - h5py==3.10.0
+      - hf-transfer==0.1.5
+      - hjson==3.1.0
+      - httpcore==1.0.2
+      - httpx==0.26.0
+      - huggingface-hub==0.22.2
+      - idna==3.6
+      - ipdb==0.13.13
+      - ipykernel==6.29.2
+      - ipython==8.21.0
+      - isoduration==20.11.0
+      - jedi==0.19.1
+      - jiwer==3.0.3
+      - joblib==1.3.2
+      - json5==0.9.14
+      - jsonpointer==2.4
+      - jsonschema==4.21.1
+      - jsonschema-specifications==2023.12.1
+      - jupyter-client==8.6.0
+      - jupyter-core==5.7.1
+      - jupyter-events==0.9.0
+      - jupyter-lsp==2.2.2
+      - jupyter-server==2.12.5
+      - jupyter-server-terminals==0.5.2
+      - jupyterlab==4.1.1
+      - jupyterlab-pygments==0.3.0
+      - jupyterlab-server==2.25.2
+      - keras==2.15.0
+      - lazy-loader==0.3
+      - libclang==16.0.6
+      - librosa==0.10.1
+      - llvmlite==0.41.1
+      - markdown==3.5.1
+      - markdown-it-py==3.0.0
+      - matplotlib-inline==0.1.6
+      - mdurl==0.1.2
+      - mistune==3.0.2
+      - ml-dtypes==0.2.0
+      - msgpack==1.0.7
+      - multidict==6.0.4
+      - multiprocess==0.70.15
+      - nbclient==0.9.0
+      - nbconvert==7.16.0
+      - nbformat==5.9.2
+      - nest-asyncio==1.6.0
+      - ninja==1.11.1.1
+      - nltk==3.8.1
+      - notebook-shim==0.2.3
+      - numba==0.58.1
+      - nvidia-cublas-cu12==12.1.3.1
+      - nvidia-cuda-cupti-cu12==12.1.105
+      - nvidia-cuda-nvrtc-cu12==12.1.105
+      - nvidia-cuda-runtime-cu12==12.1.105
+      - nvidia-cudnn-cu12==8.9.2.26
+      - nvidia-cufft-cu12==11.0.2.54
+      - nvidia-curand-cu12==10.3.2.106
+      - nvidia-cusolver-cu12==11.4.5.107
+      - nvidia-cusparse-cu12==12.1.0.106
+      - nvidia-nccl-cu12==2.20.5
+      - nvidia-nvjitlink-cu12==12.3.101
+      - nvidia-nvtx-cu12==12.1.105
+      - oauthlib==3.2.2
+      - opt-einsum==3.3.0
+      - overrides==7.7.0
+      - packaging==23.2
+      - pandas==2.1.4
+      - pandocfilters==1.5.1
+      - parso==0.8.3
+      - peft==0.7.1
+      - pexpect==4.9.0
+      - pillow==10.2.0
+      - pip==24.0
+      - platformdirs==4.1.0
+      - pooch==1.8.0
+      - prometheus-client==0.19.0
+      - prompt-toolkit==3.0.43
+      - protobuf==3.20.2
+      - psutil==5.9.7
+      - ptyprocess==0.7.0
+      - pure-eval==0.2.2
+      - py-cpuinfo==9.0.0
+      - pyarrow==14.0.2
+      - pyarrow-hotfix==0.6
+      - pyasn1==0.5.1
+      - pyasn1-modules==0.3.0
+      - pydantic==2.6.0
+      - pydantic-core==2.16.1
+      - pygments==2.17.2
+      - pynvml==11.5.0
+      - python-dateutil==2.8.2
+      - python-json-logger==2.0.7
+      - pytorch-triton==3.0.0+989adb9a29
+      - pytz==2023.3.post1
+      - pyzmq==25.1.2
+      - rapidfuzz==3.6.1
+      - referencing==0.33.0
+      - regex==2023.12.25
+      - requests-oauthlib==1.3.1
+      - responses==0.18.0
+      - rfc3339-validator==0.1.4
+      - rfc3986-validator==0.1.1
+      - rich==13.7.0
+      - rpds-py==0.17.1
+      - rsa==4.9
+      - safetensors==0.4.1
+      - scikit-learn==1.3.2
+      - scipy==1.11.4
+      - send2trash==1.8.2
+      - sentencepiece==0.1.99
+      - sentry-sdk==1.39.1
+      - setproctitle==1.3.3
+      - shtab==1.6.5
+      - six==1.16.0
+      - smmap==5.0.1
+      - sniffio==1.3.0
+      - soundfile==0.12.1
+      - soupsieve==2.5
+      - soxr==0.3.7
+      - stack-data==0.6.3
+      - tensorboard==2.15.1
+      - tensorboard-data-server==0.7.2
+      - tensorflow-cpu==2.15.0.post1
+      - tensorflow-estimator==2.15.0
+      - tensorflow-io-gcs-filesystem==0.35.0
+      - termcolor==2.4.0
+      - terminado==0.18.0
+      - threadpoolctl==3.2.0
+      - tinycss2==1.2.1
+      - tokenizers==0.15.0
+      - torch==2.4.0.dev20240323+cu121
+      - torchaudio==2.2.0.dev20240323+cu121
+      - torchvision==0.19.0.dev20240323+cu121
+      - tornado==6.4
+      - tqdm==4.66.1
+      - traitlets==5.14.1
+      - transformers==4.39.0.dev0
+      - triton==2.2.0
+      - trl==0.7.7
+      - types-python-dateutil==2.8.19.20240106
+      - typing-extensions==4.9.0
+      - tyro==0.7.0
+      - tzdata==2023.3
+      - uri-template==1.3.0
+      - urllib3==2.1.0
+      - wandb==0.16.1
+      - wcwidth==0.2.13
+      - webcolors==1.13
+      - webencodings==0.5.1
+      - websocket-client==1.7.0
+      - werkzeug==3.0.1
+      - wrapt==1.14.1
+      - xxhash==3.4.1
+      - yarl==1.9.4
+prefix: /fsx/sanchit/miniconda3/envs/venv

wandb/run-20240416_172306-uygw9yfk/files/config.yaml

@@ -0,0 +1,39 @@
+wandb_version: 1
+
+_wandb:
+  desc: null
+  value:
+    python_version: 3.11.5
+    cli_version: 0.16.1
+    framework: huggingface
+    huggingface_version: 4.40.0.dev0
+    is_jupyter_run: false
+    is_kaggle_kernel: false
+    start_time: 1713288186.348831
+    t:
+      1:
+      - 1
+      - 11
+      - 49
+      - 51
+      - 55
+      - 71
+      - 98
+      2:
+      - 1
+      - 2
+      - 3
+      - 11
+      - 49
+      - 51
+      - 55
+      - 71
+      - 98
+      3:
+      - 23
+      4: 3.11.5
+      5: 0.16.1
+      6: 4.40.0.dev0
+      8:
+      - 5
+      13: linux-x86_64

wandb/run-20240416_172306-uygw9yfk/files/output.log

@@ -0,0 +1,2124 @@
+04/16/2024 17:23:10 - WARNING - __main__ - Process rank: 0, device: cuda:0, n_gpu: 1, distributed training: True, 16-bits training: False
+04/16/2024 17:23:10 - INFO - __main__ - Training/evaluation parameters DistillationTrainingArguments(
+_n_gpu=1,
+accelerator_config={'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'gradient_accumulation_kwargs': None},
+adafactor=False,
+adam_beta1=0.9,
+adam_beta2=0.999,
+adam_epsilon=1e-08,
+auto_find_batch_size=False,
+bf16=False,
+bf16_full_eval=False,
+data_seed=None,
+dataloader_drop_last=False,
+dataloader_num_workers=4,
+dataloader_persistent_workers=False,
+dataloader_pin_memory=True,
+dataloader_prefetch_factor=None,
+ddp_backend=None,
+ddp_broadcast_buffers=None,
+ddp_bucket_cap_mb=None,
+ddp_find_unused_parameters=None,
+ddp_timeout=7200,
+debug=[],
+deepspeed=None,
+disable_tqdm=False,
+dispatch_batches=None,
+do_eval=True,
+do_predict=False,
+do_train=True,
+dtype=bfloat16,
+eval_accumulation_steps=None,
+eval_delay=0,
+eval_steps=None,
+evaluation_strategy=IntervalStrategy.EPOCH,
+fp16=False,
+fp16_backend=auto,
+fp16_full_eval=False,
+fp16_opt_level=O1,
+freeze_lm_head=False,
+fsdp=[],
+fsdp_config={'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False},
+fsdp_min_num_params=0,
+fsdp_transformer_layer_cls_to_wrap=None,
+full_determinism=False,
+generation_config=None,
+generation_max_length=None,
+generation_num_beams=None,
+gradient_accumulation_steps=1,
+gradient_checkpointing=False,
+gradient_checkpointing_kwargs=None,
+greater_is_better=None,
+group_by_length=False,
+half_precision_backend=auto,
+hub_always_push=False,
+hub_model_id=None,
+hub_private_repo=False,
+hub_strategy=HubStrategy.EVERY_SAVE,
+hub_token=<HUB_TOKEN>,
+ignore_data_skip=False,
+include_inputs_for_metrics=False,
+include_num_input_tokens_seen=False,
+include_tokens_per_second=False,
+jit_mode_eval=False,
+kl_weight=1.0,
+label_names=None,
+label_smoothing_factor=0.0,
+learning_rate=0.0003,
+length_column_name=length,
+load_best_model_at_end=False,
+local_rank=0,
+log_level=passive,
+log_level_replica=warning,
+log_on_each_node=True,
+logging_dir=./runs/Apr16_17-23-01_ip-26-0-164-187,
+logging_first_step=False,
+logging_nan_inf_filter=True,
+logging_steps=25,
+logging_strategy=IntervalStrategy.STEPS,
+lr_scheduler_kwargs={},
+lr_scheduler_type=SchedulerType.LINEAR,
+max_grad_norm=1.0,
+max_steps=50000,
+metric_for_best_model=None,
+mp_parameters=,
+neftune_noise_alpha=None,
+no_cuda=False,
+num_train_epochs=3.0,
+optim=OptimizerNames.ADAMW_TORCH,
+optim_args=None,
+optim_target_modules=None,
+output_dir=./,
+output_router_logits=True,
+overwrite_output_dir=True,
+past_index=-1,
+per_device_eval_batch_size=8,
+per_device_train_batch_size=8,
+predict_with_generate=False,
+prediction_loss_only=False,
+push_to_hub=False,
+push_to_hub_model_id=None,
+push_to_hub_organization=None,
+push_to_hub_token=<PUSH_TO_HUB_TOKEN>,
+ray_scope=last,
+remove_unused_columns=True,
+report_to=['wandb'],
+resume_from_checkpoint=None,
+run_name=./,
+save_on_each_node=False,
+save_only_model=False,
+save_safetensors=True,
+save_steps=500,
+save_strategy=IntervalStrategy.EPOCH,
+save_total_limit=1,
+seed=42,
+skip_memory_metrics=True,
+sortish_sampler=False,
+split_batches=None,
+temperature=2.0,
+tf32=None,
+torch_compile=False,
+torch_compile_backend=None,
+torch_compile_mode=None,
+torchdynamo=None,
+tpu_metrics_debug=False,
+tpu_num_cores=None,
+use_cpu=False,
+use_ipex=False,
+use_legacy_prediction_loop=False,
+use_mps_device=False,
+warmup_ratio=0.0,
+warmup_steps=500,
+weight_decay=0.0,
+)
+Combining datasets...:   0%|                                                                      | 0/8 [00:00<?, ?it/s]
+Resolving data files: 100%|█████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 101885.93it/s]
+Loading dataset shards: 100%|█████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 2253.19it/s]
+Resolving data files: 100%|██████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 94019.27it/s]
+
+Resolving data files: 100%|████████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 138.64it/s]
+Resolving data files: 100%|██████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 97668.14it/s]
+Resolving data files: 100%|████████████████████████████████████���█████████████████████| 18/18 [00:00<00:00, 96175.12it/s]
+Resolving data files: 100%|████████████████████████████████████████████████████████████| 43/43 [00:00<00:00, 414.13it/s]
+Loading dataset shards: 100%|█████████████████████████████████████████████████████████| 43/43 [00:00<00:00, 2036.64it/s]
+Resolving data files: 100%|████████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 152.49it/s]
+Resolving data files: 100%|███████████████████████████████████████████████████████| 139/139 [00:00<00:00, 322460.32it/s]
+Loading dataset shards: 100%|███████████████████████████████████████████████████████| 138/138 [00:00<00:00, 2441.33it/s]
+Resolving data files: 100%|████████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 158.30it/s]
+Resolving data files:  89%|█████████████████████████████████████████████████████▎      | 16/18 [00:00<00:00, 141.58it/s]
+Loading dataset shards: 100%|███████████████████████████████████████████████████████| 117/117 [00:00<00:00, 2419.04it/s]
+Combining datasets...:  88%|██████████████████████████████████████████████████████▎       | 7/8 [00:24<00:03,  3.60s/it]
+
+Combining datasets...: 100%|██████████████████████████████████████████████████████████████| 8/8 [00:30<00:00,  3.82s/it]
+Resolving data files: 100%|█████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 104134.44it/s]
+Resolving data files: 100%|██████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 98176.17it/s]
+Resolving data files: 100%|██████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 93553.25it/s]
+Resolving data files: 100%|██████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 90851.35it/s]
+Resolving data files: 100%|██████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 98560.67it/s]
+Resolving data files: 100%|██████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 99996.65it/s]
+Resolving data files: 100%|█████████████████████████████████████████████████████████| 43/43 [00:00<00:00, 183287.67it/s]
+Resolving data files: 100%|██████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 95325.09it/s]
+Resolving data files: 100%|███████████████████████████████████████████████████████| 139/139 [00:00<00:00, 366902.62it/s]
+Resolving data files: 100%|█████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 103705.32it/s]
+Resolving data files: 100%|███████████████████████████████████████████████████████| 118/118 [00:00<00:00, 308174.27it/s]
+Resolving data files: 100%|███████��██████████████████████████████████████████████████| 18/18 [00:00<00:00, 85211.59it/s]
+loading configuration file config.json from cache at /fsx/sanchit/cache/models--sanchit-gandhi--Mistral-7B-v0.1-6-layer/snapshots/d4e2300e8038196385e9106614a4d7b6c5b70211/config.json
+Model config MistralConfig {
+  "_name_or_path": "sanchit-gandhi/Mistral-7B-v0.1-6-layer",
+  "architectures": [
+    "MistralForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 14336,
+  "max_position_embeddings": 32768,
+  "model_type": "mistral",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 6,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "sliding_window": 4096,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.0.dev0",
+  "use_cache": true,
+  "vocab_size": 32000
+}
+loading file tokenizer.model from cache at /fsx/sanchit/cache/models--sanchit-gandhi--Mistral-7B-v0.1-6-layer/snapshots/d4e2300e8038196385e9106614a4d7b6c5b70211/tokenizer.model
+loading file tokenizer.json from cache at /fsx/sanchit/cache/models--sanchit-gandhi--Mistral-7B-v0.1-6-layer/snapshots/d4e2300e8038196385e9106614a4d7b6c5b70211/tokenizer.json
+loading file added_tokens.json from cache at None
+loading file special_tokens_map.json from cache at /fsx/sanchit/cache/models--sanchit-gandhi--Mistral-7B-v0.1-6-layer/snapshots/d4e2300e8038196385e9106614a4d7b6c5b70211/special_tokens_map.json
+loading file tokenizer_config.json from cache at /fsx/sanchit/cache/models--sanchit-gandhi--Mistral-7B-v0.1-6-layer/snapshots/d4e2300e8038196385e9106614a4d7b6c5b70211/tokenizer_config.json
+loading configuration file config.json from cache at /fsx/sanchit/cache/models--mistralai--Mistral-7B-v0.1/snapshots/26bca36bde8333b5d7f72e9ed20ccda6a618af24/config.json
+Model config MistralConfig {
+  "_name_or_path": "mistralai/Mistral-7B-v0.1",
+  "architectures": [
+    "MistralForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 14336,
+  "max_position_embeddings": 32768,
+  "model_type": "mistral",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "sliding_window": 4096,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.40.0.dev0",
+  "use_cache": true,
+  "vocab_size": 32000
+}
+The device_map was not initialized. Setting device_map to {'':torch.cuda.current_device()}. If you want to use the model for inference, please set device_map ='auto'
+loading weights file model.safetensors from cache at /fsx/sanchit/cache/models--mistralai--Mistral-7B-v0.1/snapshots/26bca36bde8333b5d7f72e9ed20ccda6a618af24/model.safetensors.index.json
+Downloading shards: 100%|███████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 1479.73it/s]
+Instantiating MistralForCausalLM model under default dtype torch.bfloat16.
+Generate config GenerationConfig {
+  "bos_token_id": 1,
+  "eos_token_id": 2
+}
+
+Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████| 2/2 [00:05<00:00,  2.99s/it]
+All model checkpoint weights were used when initializing MistralForCausalLM.
+All the weights of MistralForCausalLM were initialized from the model checkpoint at mistralai/Mistral-7B-v0.1.
+If your task is similar to the task the model of the checkpoint was trained on, you can already use MistralForCausalLM for predictions without further training.
+loading configuration file generation_config.json from cache at /fsx/sanchit/cache/models--mistralai--Mistral-7B-v0.1/snapshots/26bca36bde8333b5d7f72e9ed20ccda6a618af24/generation_config.json
+Generate config GenerationConfig {
+  "bos_token_id": 1,
+  "eos_token_id": 2
+}
+loading weights file model.safetensors from cache at /fsx/sanchit/cache/models--sanchit-gandhi--Mistral-7B-v0.1-6-layer/snapshots/d4e2300e8038196385e9106614a4d7b6c5b70211/model.safetensors.index.json
+Downloading shards: 100%|███████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 1308.47it/s]
+Generate config GenerationConfig {
+  "bos_token_id": 1,
+  "eos_token_id": 2
+}
+Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  4.66it/s]
+All model checkpoint weights were used when initializing MistralForCausalLM.
+All the weights of MistralForCausalLM were initialized from the model checkpoint at sanchit-gandhi/Mistral-7B-v0.1-6-layer.
+If your task is similar to the task the model of the checkpoint was trained on, you can already use MistralForCausalLM for predictions without further training.
+loading configuration file generation_config.json from cache at /fsx/sanchit/cache/models--sanchit-gandhi--Mistral-7B-v0.1-6-layer/snapshots/d4e2300e8038196385e9106614a4d7b6c5b70211/generation_config.json
+Generate config GenerationConfig {
+  "bos_token_id": 1,
+  "eos_token_id": 2
+}
+tokenizer config file saved in ./tokenizer_config.json
+Special tokens file saved in ./special_tokens_map.json
+Configuration saved in ./config.json
+Configuration saved in ./generation_config.json
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+preprocess train dataset (num_proc=32): 100%|██████████████████████| 31056744/31056744 [54:46<00:00, 9450.26 examples/s]
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+filtering train dataset (num_proc=32): 100%|██████████████████████| 31056744/31056744 [05:21<00:00, 96735.06 examples/s]
+04/16/2024 18:24:47 - INFO - __main__ - max_steps is given, it will override any value given in num_train_epochs
+04/16/2024 18:25:15 - INFO - __main__ - ***** Running training *****
+04/16/2024 18:25:15 - INFO - __main__ -   Num examples = 3200000
+04/16/2024 18:25:15 - INFO - __main__ -   Num epochs = 1
+04/16/2024 18:25:15 - INFO - __main__ -   Instantaneous batch size per device = 8
+04/16/2024 18:25:15 - INFO - __main__ -   Gradient accumulation steps = 1
+04/16/2024 18:25:15 - INFO - __main__ -   Total train batch size (w. parallel & distributed) = 64
+04/16/2024 18:25:15 - INFO - __main__ -   Total optimization steps = 50000
+
+Train steps ... :   0%|                                                           | 1/50000 [00:10<140:57:09, 10.15s/it]Traceback (most recent call last):
+  File "/fsx/sanchit/mistral-debug-4bit/run_distillation.py", line 1474, in <module>
+    main()
+  File "/fsx/sanchit/mistral-debug-4bit/run_distillation.py", line 1326, in main
+    loss, train_metric = train_step(batch, temperature=training_args.temperature)
+                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/fsx/sanchit/mistral-debug-4bit/run_distillation.py", line 1211, in train_step
+    kl_loss = kl_divergence(teacher_distribution, student_distribution, batch["labels"]) * temperature**2
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/fsx/sanchit/mistral-debug-4bit/run_distillation.py", line 1183, in kl_divergence
+    divergence = kl_loss(log_predicted_distribution, target_distribution)
+                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/fsx/sanchit/miniconda3/envs/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/fsx/sanchit/miniconda3/envs/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1536, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/fsx/sanchit/miniconda3/envs/venv/lib/python3.11/site-packages/torch/nn/modules/loss.py", line 470, in forward
+    return F.kl_div(input, target, reduction=self.reduction, log_target=self.log_target)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/fsx/sanchit/miniconda3/envs/venv/lib/python3.11/site-packages/torch/nn/functional.py", line 2990, in kl_div
+    reduced = torch.kl_div(input, target, reduction_enum, log_target=log_target)
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.95 GiB. GPU
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/fsx/sanchit/mistral-debug-4bit/run_distillation.py", line 1474, in <module>
+[rank0]:     main()
+[rank0]:   File "/fsx/sanchit/mistral-debug-4bit/run_distillation.py", line 1326, in main
+[rank0]:     loss, train_metric = train_step(batch, temperature=training_args.temperature)
+[rank0]:                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/fsx/sanchit/mistral-debug-4bit/run_distillation.py", line 1211, in train_step
+[rank0]:     kl_loss = kl_divergence(teacher_distribution, student_distribution, batch["labels"]) * temperature**2
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/fsx/sanchit/mistral-debug-4bit/run_distillation.py", line 1183, in kl_divergence
+[rank0]:     divergence = kl_loss(log_predicted_distribution, target_distribution)
+[rank0]:                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/fsx/sanchit/miniconda3/envs/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/fsx/sanchit/miniconda3/envs/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1536, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/fsx/sanchit/miniconda3/envs/venv/lib/python3.11/site-packages/torch/nn/modules/loss.py", line 470, in forward
+[rank0]:     return F.kl_div(input, target, reduction=self.reduction, log_target=self.log_target)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/fsx/sanchit/miniconda3/envs/venv/lib/python3.11/site-packages/torch/nn/functional.py", line 2990, in kl_div
+[rank0]:     reduced = torch.kl_div(input, target, reduction_enum, log_target=log_target)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]: torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.95 GiB. GPU

wandb/run-20240416_172306-uygw9yfk/files/requirements.txt

@@ -0,0 +1,223 @@
+absl-py==2.0.0
+accelerate==0.29.2
+aiohttp==3.9.1
+aiosignal==1.3.1
+annotated-types==0.6.0
+anyio==4.2.0
+appdirs==1.4.4
+argon2-cffi-bindings==21.2.0
+argon2-cffi==23.1.0
+arrow==1.3.0
+asttokens==2.4.1
+astunparse==1.6.3
+async-lru==2.0.4
+attrs==23.1.0
+audioread==3.0.1
+babel==2.14.0
+beautifulsoup4==4.12.3
+bitsandbytes==0.43.1
+bleach==6.1.0
+brotli==1.0.9
+cachetools==5.3.2
+certifi==2023.11.17
+cffi==1.16.0
+chardet==5.2.0
+charset-normalizer==2.0.4
+click==8.1.7
+comm==0.2.1
+cryptography==41.0.7
+datasets==2.18.1.dev0
+debugpy==1.8.1
+decorator==5.1.1
+deepspeed==0.12.2
+defusedxml==0.7.1
+dill==0.3.7
+docker-pycreds==0.4.0
+docstring-parser==0.15
+einops==0.7.0
+evaluate==0.4.0
+executing==2.0.1
+fastjsonschema==2.19.1
+filelock==3.13.1
+flatbuffers==23.5.26
+fqdn==1.5.1
+frozenlist==1.4.1
+fsspec==2023.10.0
+gast==0.5.4
+gitdb==4.0.11
+gitpython==3.1.40
+gmpy2==2.1.2
+google-auth-oauthlib==1.2.0
+google-auth==2.26.1
+google-pasta==0.2.0
+grpcio==1.60.0
+h11==0.14.0
+h5py==3.10.0
+hf-transfer==0.1.5
+hjson==3.1.0
+httpcore==1.0.2
+httpx==0.26.0
+huggingface-hub==0.22.2
+idna==3.4
+ipdb==0.13.13
+ipykernel==6.29.2
+ipython==8.21.0
+isoduration==20.11.0
+jedi==0.19.1
+jinja2==3.1.2
+jiwer==3.0.3
+joblib==1.3.2
+json5==0.9.14
+jsonpointer==2.4
+jsonschema-specifications==2023.12.1
+jsonschema==4.21.1
+jupyter-client==8.6.0
+jupyter-core==5.7.1
+jupyter-events==0.9.0
+jupyter-lsp==2.2.2
+jupyter-server-terminals==0.5.2
+jupyter-server==2.12.5
+jupyterlab-pygments==0.3.0
+jupyterlab-server==2.25.2
+jupyterlab==4.1.1
+keras==2.15.0
+lazy-loader==0.3
+libclang==16.0.6
+librosa==0.10.1
+llvmlite==0.41.1
+markdown-it-py==3.0.0
+markdown==3.5.1
+markupsafe==2.1.1
+matplotlib-inline==0.1.6
+mdurl==0.1.2
+mistune==3.0.2
+mkl-fft==1.3.8
+mkl-random==1.2.4
+mkl-service==2.4.0
+ml-dtypes==0.2.0
+mpmath==1.3.0
+msgpack==1.0.7
+multidict==6.0.4
+multiprocess==0.70.15
+nbclient==0.9.0
+nbconvert==7.16.0
+nbformat==5.9.2
+nest-asyncio==1.6.0
+networkx==3.1
+ninja==1.11.1.1
+nltk==3.8.1
+notebook-shim==0.2.3
+numba==0.58.1
+numpy==1.26.2
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.20.5
+nvidia-nvjitlink-cu12==12.3.101
+nvidia-nvtx-cu12==12.1.105
+oauthlib==3.2.2
+opt-einsum==3.3.0
+overrides==7.7.0
+packaging==23.2
+pandas==2.1.4
+pandocfilters==1.5.1
+parso==0.8.3
+peft==0.7.1
+pexpect==4.9.0
+pillow==10.2.0
+pip==24.0
+platformdirs==4.1.0
+pooch==1.8.0
+prometheus-client==0.19.0
+prompt-toolkit==3.0.43
+protobuf==3.20.2
+psutil==5.9.7
+ptyprocess==0.7.0
+pure-eval==0.2.2
+py-cpuinfo==9.0.0
+pyarrow-hotfix==0.6
+pyarrow==14.0.2
+pyasn1-modules==0.3.0
+pyasn1==0.5.1
+pycparser==2.21
+pydantic-core==2.16.1
+pydantic==2.6.0
+pygments==2.17.2
+pynvml==11.5.0
+pyopenssl==23.2.0
+pysocks==1.7.1
+python-dateutil==2.8.2
+python-json-logger==2.0.7
+pytorch-triton==3.0.0+989adb9a29
+pytz==2023.3.post1
+pyyaml==6.0.1
+pyzmq==25.1.2
+rapidfuzz==3.6.1
+referencing==0.33.0
+regex==2023.12.25
+requests-oauthlib==1.3.1
+requests==2.31.0
+responses==0.18.0
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rich==13.7.0
+rpds-py==0.17.1
+rsa==4.9
+safetensors==0.4.1
+scikit-learn==1.3.2
+scipy==1.11.4
+send2trash==1.8.2
+sentencepiece==0.1.99
+sentry-sdk==1.39.1
+setproctitle==1.3.3
+setuptools==68.2.2
+shtab==1.6.5
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.0
+soundfile==0.12.1
+soupsieve==2.5
+soxr==0.3.7
+stack-data==0.6.3
+sympy==1.12
+tensorboard-data-server==0.7.2
+tensorboard==2.15.1
+tensorflow-cpu==2.15.0.post1
+tensorflow-estimator==2.15.0
+tensorflow-io-gcs-filesystem==0.35.0
+termcolor==2.4.0
+terminado==0.18.0
+threadpoolctl==3.2.0
+tinycss2==1.2.1
+tokenizers==0.15.0
+torch==2.4.0.dev20240323+cu121
+torchaudio==2.2.0.dev20240323+cu121
+torchvision==0.19.0.dev20240323+cu121
+tornado==6.4
+tqdm==4.66.1
+traitlets==5.14.1
+transformers==4.39.0.dev0
+triton==2.2.0
+trl==0.7.7
+types-python-dateutil==2.8.19.20240106
+typing-extensions==4.10.0
+tyro==0.7.0
+tzdata==2023.3
+uri-template==1.3.0
+urllib3==1.26.18
+wandb==0.16.1
+wcwidth==0.2.13
+webcolors==1.13
+webencodings==0.5.1
+websocket-client==1.7.0
+werkzeug==3.0.1
+wheel==0.41.2
+wrapt==1.14.1
+xxhash==3.4.1
+yarl==1.9.4

wandb/run-20240416_172306-uygw9yfk/files/wandb-metadata.json

@@ -0,0 +1,552 @@
+{
+    "os": "Linux-5.15.0-1048-aws-x86_64-with-glibc2.31",
+    "python": "3.11.5",
+    "heartbeatAt": "2024-04-16T17:23:06.717620",
+    "startedAt": "2024-04-16T17:23:06.329363",
+    "docker": null,
+    "cuda": null,
+    "args": [
+        "config_mistral.yaml"
+    ],
+    "state": "running",
+    "program": "/fsx/sanchit/mistral-debug-4bit/run_distillation.py",
+    "codePathLocal": "run_distillation.py",
+    "codePath": "run_distillation.py",
+    "host": "ip-26-0-164-187",
+    "username": "sanchit",
+    "executable": "/fsx/sanchit/miniconda3/envs/venv/bin/python",
+    "cpu_count": 96,
+    "cpu_count_logical": 96,
+    "cpu_freq": {
+        "current": 2659.601562499998,
+        "min": 0.0,
+        "max": 0.0
+    },
+    "cpu_freq_per_core": [
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3571.94,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        }
+    ],
+    "disk": {
+        "/": {
+            "total": 290.7472343444824,
+            "used": 58.41170883178711
+        }
+    },
+    "gpu": "NVIDIA H100 80GB HBM3",
+    "gpu_count": 8,
+    "gpu_devices": [
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        }
+    ],
+    "memory": {
+        "total": 1999.9855461120605
+    }
+}

wandb/run-20240416_172306-uygw9yfk/files/wandb-summary.json

@@ -0,0 +1 @@
+{"_wandb": {"runtime": 3740}}

wandb/run-20240416_172306-uygw9yfk/logs/debug.log

@@ -0,0 +1,28 @@
+2024-04-16 17:23:06,340 INFO    MainThread:940676 [wandb_setup.py:_flush():76] Current SDK version is 0.16.1
+2024-04-16 17:23:06,340 INFO    MainThread:940676 [wandb_setup.py:_flush():76] Configure stats pid to 940676
+2024-04-16 17:23:06,340 INFO    MainThread:940676 [wandb_setup.py:_flush():76] Loading settings from /admin/home/sanchit/.config/wandb/settings
+2024-04-16 17:23:06,340 INFO    MainThread:940676 [wandb_setup.py:_flush():76] Loading settings from /fsx/sanchit/mistral-debug-4bit/wandb/settings
+2024-04-16 17:23:06,340 INFO    MainThread:940676 [wandb_setup.py:_flush():76] Loading settings from environment variables: {}
+2024-04-16 17:23:06,341 INFO    MainThread:940676 [wandb_setup.py:_flush():76] Applying setup settings: {'_disable_service': False}
+2024-04-16 17:23:06,341 INFO    MainThread:940676 [wandb_setup.py:_flush():76] Inferring run settings from compute environment: {'program_relpath': 'run_distillation.py', 'program_abspath': '/fsx/sanchit/mistral-debug-4bit/run_distillation.py', 'program': '/fsx/sanchit/mistral-debug-4bit/run_distillation.py'}
+2024-04-16 17:23:06,341 INFO    MainThread:940676 [wandb_init.py:_log_setup():524] Logging user logs to /fsx/sanchit/mistral-debug-4bit/wandb/run-20240416_172306-uygw9yfk/logs/debug.log
+2024-04-16 17:23:06,341 INFO    MainThread:940676 [wandb_init.py:_log_setup():525] Logging internal logs to /fsx/sanchit/mistral-debug-4bit/wandb/run-20240416_172306-uygw9yfk/logs/debug-internal.log
+2024-04-16 17:23:06,341 INFO    MainThread:940676 [wandb_init.py:init():564] calling init triggers
+2024-04-16 17:23:06,341 INFO    MainThread:940676 [wandb_init.py:init():571] wandb.init called with sweep_config: {}
+config: {}
+2024-04-16 17:23:06,341 INFO    MainThread:940676 [wandb_init.py:init():614] starting backend
+2024-04-16 17:23:06,341 INFO    MainThread:940676 [wandb_init.py:init():618] setting up manager
+2024-04-16 17:23:06,346 INFO    MainThread:940676 [backend.py:_multiprocessing_setup():105] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-04-16 17:23:06,348 INFO    MainThread:940676 [wandb_init.py:init():624] backend started and connected
+2024-04-16 17:23:06,353 INFO    MainThread:940676 [wandb_init.py:init():716] updated telemetry
+2024-04-16 17:23:06,355 INFO    MainThread:940676 [wandb_init.py:init():749] communicating run to backend with 90.0 second timeout
+2024-04-16 17:23:06,543 INFO    MainThread:940676 [wandb_run.py:_on_init():2254] communicating current version
+2024-04-16 17:23:06,597 INFO    MainThread:940676 [wandb_run.py:_on_init():2263] got version response upgrade_message: "wandb version 0.16.6 is available!  To upgrade, please run:\n $ pip install wandb --upgrade"
+
+2024-04-16 17:23:06,597 INFO    MainThread:940676 [wandb_init.py:init():800] starting run threads in backend
+2024-04-16 17:23:10,890 INFO    MainThread:940676 [wandb_run.py:_console_start():2233] atexit reg
+2024-04-16 17:23:10,890 INFO    MainThread:940676 [wandb_run.py:_redirect():2088] redirect: wrap_raw
+2024-04-16 17:23:10,890 INFO    MainThread:940676 [wandb_run.py:_redirect():2153] Wrapping output streams.
+2024-04-16 17:23:10,890 INFO    MainThread:940676 [wandb_run.py:_redirect():2178] Redirects installed.
+2024-04-16 17:23:10,891 INFO    MainThread:940676 [wandb_init.py:init():841] run started, returning control to user process
+2024-04-16 18:25:33,477 WARNING MsgRouterThr:940676 [router.py:message_loop():77] message_loop has been closed

wandb/run-20240416_172306-uygw9yfk/run-uygw9yfk.wandb

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

wandb/run-20240416_205309-xdytsc71/files/conda-environment.yaml

@@ -0,0 +1,300 @@
+name: venv
+channels:
+  - pytorch
+  - nvidia
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - blas=1.0=mkl
+  - brotli-python=1.0.9=py311h6a678d5_7
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2023.12.12=h06a4308_0
+  - certifi=2023.11.17=py311h06a4308_0
+  - cffi=1.16.0=py311h5eee18b_0
+  - cryptography=41.0.7=py311hdda0065_0
+  - cuda-cudart=12.1.105=0
+  - cuda-cupti=12.1.105=0
+  - cuda-libraries=12.1.0=0
+  - cuda-nvrtc=12.1.105=0
+  - cuda-nvtx=12.1.105=0
+  - cuda-opencl=12.3.101=0
+  - cuda-runtime=12.1.0=0
+  - ffmpeg=4.3=hf484d3e_0
+  - filelock=3.13.1=py311h06a4308_0
+  - freetype=2.12.1=h4a9f257_0
+  - giflib=5.2.1=h5eee18b_3
+  - gmp=6.2.1=h295c915_3
+  - gmpy2=2.1.2=py311hc9b5ff0_0
+  - gnutls=3.6.15=he1e5248_0
+  - intel-openmp=2023.1.0=hdb19cb5_46306
+  - jinja2=3.1.2=py311h06a4308_0
+  - jpeg=9e=h5eee18b_1
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - lerc=3.0=h295c915_0
+  - libcublas=12.1.0.26=0
+  - libcufft=11.0.2.4=0
+  - libcufile=1.8.1.2=0
+  - libcurand=10.3.4.101=0
+  - libcusolver=11.4.4.55=0
+  - libcusparse=12.0.2.55=0
+  - libdeflate=1.17=h5eee18b_1
+  - libffi=3.4.4=h6a678d5_0
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgomp=11.2.0=h1234567_1
+  - libiconv=1.16=h7f8727e_2
+  - libidn2=2.3.4=h5eee18b_0
+  - libjpeg-turbo=2.0.0=h9bf148f_0
+  - libnpp=12.0.2.50=0
+  - libnvjitlink=12.1.105=0
+  - libnvjpeg=12.1.1.14=0
+  - libpng=1.6.39=h5eee18b_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libtasn1=4.19.0=h5eee18b_0
+  - libtiff=4.5.1=h6a678d5_0
+  - libunistring=0.9.10=h27cfd23_0
+  - libuuid=1.41.5=h5eee18b_0
+  - libwebp=1.3.2=h11a3e52_0
+  - libwebp-base=1.3.2=h5eee18b_0
+  - llvm-openmp=14.0.6=h9e868ea_0
+  - lz4-c=1.9.4=h6a678d5_0
+  - markupsafe=2.1.1=py311h5eee18b_0
+  - mkl=2023.1.0=h213fc3f_46344
+  - mkl-service=2.4.0=py311h5eee18b_1
+  - mkl_fft=1.3.8=py311h5eee18b_0
+  - mkl_random=1.2.4=py311hdb19cb5_0
+  - mpc=1.1.0=h10f8cd9_1
+  - mpfr=4.0.2=hb69a4c5_1
+  - mpmath=1.3.0=py311h06a4308_0
+  - ncurses=6.4=h6a678d5_0
+  - nettle=3.7.3=hbbd107a_1
+  - networkx=3.1=py311h06a4308_0
+  - numpy=1.26.2=py311h08b1b3b_0
+  - numpy-base=1.26.2=py311hf175353_0
+  - openh264=2.1.1=h4ff587b_0
+  - openjpeg=2.4.0=h3ad879b_0
+  - openssl=3.0.12=h7f8727e_0
+  - pycparser=2.21=pyhd3eb1b0_0
+  - pyopenssl=23.2.0=py311h06a4308_0
+  - pysocks=1.7.1=py311h06a4308_0
+  - python=3.11.5=h955ad1f_0
+  - pytorch-cuda=12.1=ha16c6d3_5
+  - pytorch-mutex=1.0=cuda
+  - pyyaml=6.0.1=py311h5eee18b_0
+  - readline=8.2=h5eee18b_0
+  - requests=2.31.0=py311h06a4308_0
+  - setuptools=68.2.2=py311h06a4308_0
+  - sqlite=3.41.2=h5eee18b_0
+  - sympy=1.12=py311h06a4308_0
+  - tbb=2021.8.0=hdb19cb5_0
+  - tk=8.6.12=h1ccaba5_0
+  - wheel=0.41.2=py311h06a4308_0
+  - xz=5.4.5=h5eee18b_0
+  - yaml=0.2.5=h7b6447c_0
+  - zlib=1.2.13=h5eee18b_0
+  - zstd=1.5.5=hc292b87_0
+  - pip:
+      - absl-py==2.0.0
+      - accelerate==0.29.2
+      - aiohttp==3.9.1
+      - aiosignal==1.3.1
+      - annotated-types==0.6.0
+      - anyio==4.2.0
+      - appdirs==1.4.4
+      - argon2-cffi==23.1.0
+      - argon2-cffi-bindings==21.2.0
+      - arrow==1.3.0
+      - asttokens==2.4.1
+      - astunparse==1.6.3
+      - async-lru==2.0.4
+      - attrs==23.1.0
+      - audioread==3.0.1
+      - babel==2.14.0
+      - beautifulsoup4==4.12.3
+      - bitsandbytes==0.43.1
+      - bleach==6.1.0
+      - cachetools==5.3.2
+      - chardet==5.2.0
+      - charset-normalizer==3.3.2
+      - click==8.1.7
+      - comm==0.2.1
+      - datasets==2.18.1.dev0
+      - debugpy==1.8.1
+      - decorator==5.1.1
+      - deepspeed==0.12.2
+      - defusedxml==0.7.1
+      - dill==0.3.7
+      - docker-pycreds==0.4.0
+      - docstring-parser==0.15
+      - einops==0.7.0
+      - evaluate==0.4.0
+      - executing==2.0.1
+      - fastjsonschema==2.19.1
+      - flatbuffers==23.5.26
+      - fqdn==1.5.1
+      - frozenlist==1.4.1
+      - fsspec==2023.10.0
+      - gast==0.5.4
+      - gitdb==4.0.11
+      - gitpython==3.1.40
+      - google-auth==2.26.1
+      - google-auth-oauthlib==1.2.0
+      - google-pasta==0.2.0
+      - grpcio==1.60.0
+      - h11==0.14.0
+      - h5py==3.10.0
+      - hf-transfer==0.1.5
+      - hjson==3.1.0
+      - httpcore==1.0.2
+      - httpx==0.26.0
+      - huggingface-hub==0.22.2
+      - idna==3.6
+      - ipdb==0.13.13
+      - ipykernel==6.29.2
+      - ipython==8.21.0
+      - isoduration==20.11.0
+      - jedi==0.19.1
+      - jiwer==3.0.3
+      - joblib==1.3.2
+      - json5==0.9.14
+      - jsonpointer==2.4
+      - jsonschema==4.21.1
+      - jsonschema-specifications==2023.12.1
+      - jupyter-client==8.6.0
+      - jupyter-core==5.7.1
+      - jupyter-events==0.9.0
+      - jupyter-lsp==2.2.2
+      - jupyter-server==2.12.5
+      - jupyter-server-terminals==0.5.2
+      - jupyterlab==4.1.1
+      - jupyterlab-pygments==0.3.0
+      - jupyterlab-server==2.25.2
+      - keras==2.15.0
+      - lazy-loader==0.3
+      - libclang==16.0.6
+      - librosa==0.10.1
+      - llvmlite==0.41.1
+      - markdown==3.5.1
+      - markdown-it-py==3.0.0
+      - matplotlib-inline==0.1.6
+      - mdurl==0.1.2
+      - mistune==3.0.2
+      - ml-dtypes==0.2.0
+      - msgpack==1.0.7
+      - multidict==6.0.4
+      - multiprocess==0.70.15
+      - nbclient==0.9.0
+      - nbconvert==7.16.0
+      - nbformat==5.9.2
+      - nest-asyncio==1.6.0
+      - ninja==1.11.1.1
+      - nltk==3.8.1
+      - notebook-shim==0.2.3
+      - numba==0.58.1
+      - nvidia-cublas-cu12==12.1.3.1
+      - nvidia-cuda-cupti-cu12==12.1.105
+      - nvidia-cuda-nvrtc-cu12==12.1.105
+      - nvidia-cuda-runtime-cu12==12.1.105
+      - nvidia-cudnn-cu12==8.9.2.26
+      - nvidia-cufft-cu12==11.0.2.54
+      - nvidia-curand-cu12==10.3.2.106
+      - nvidia-cusolver-cu12==11.4.5.107
+      - nvidia-cusparse-cu12==12.1.0.106
+      - nvidia-nccl-cu12==2.20.5
+      - nvidia-nvjitlink-cu12==12.3.101
+      - nvidia-nvtx-cu12==12.1.105
+      - oauthlib==3.2.2
+      - opt-einsum==3.3.0
+      - overrides==7.7.0
+      - packaging==23.2
+      - pandas==2.1.4
+      - pandocfilters==1.5.1
+      - parso==0.8.3
+      - peft==0.7.1
+      - pexpect==4.9.0
+      - pillow==10.2.0
+      - pip==24.0
+      - platformdirs==4.1.0
+      - pooch==1.8.0
+      - prometheus-client==0.19.0
+      - prompt-toolkit==3.0.43
+      - protobuf==3.20.2
+      - psutil==5.9.7
+      - ptyprocess==0.7.0
+      - pure-eval==0.2.2
+      - py-cpuinfo==9.0.0
+      - pyarrow==14.0.2
+      - pyarrow-hotfix==0.6
+      - pyasn1==0.5.1
+      - pyasn1-modules==0.3.0
+      - pydantic==2.6.0
+      - pydantic-core==2.16.1
+      - pygments==2.17.2
+      - pynvml==11.5.0
+      - python-dateutil==2.8.2
+      - python-json-logger==2.0.7
+      - pytorch-triton==3.0.0+989adb9a29
+      - pytz==2023.3.post1
+      - pyzmq==25.1.2
+      - rapidfuzz==3.6.1
+      - referencing==0.33.0
+      - regex==2023.12.25
+      - requests-oauthlib==1.3.1
+      - responses==0.18.0
+      - rfc3339-validator==0.1.4
+      - rfc3986-validator==0.1.1
+      - rich==13.7.0
+      - rpds-py==0.17.1
+      - rsa==4.9
+      - safetensors==0.4.1
+      - scikit-learn==1.3.2
+      - scipy==1.11.4
+      - send2trash==1.8.2
+      - sentencepiece==0.1.99
+      - sentry-sdk==1.39.1
+      - setproctitle==1.3.3
+      - shtab==1.6.5
+      - six==1.16.0
+      - smmap==5.0.1
+      - sniffio==1.3.0
+      - soundfile==0.12.1
+      - soupsieve==2.5
+      - soxr==0.3.7
+      - stack-data==0.6.3
+      - tensorboard==2.15.1
+      - tensorboard-data-server==0.7.2
+      - tensorflow-cpu==2.15.0.post1
+      - tensorflow-estimator==2.15.0
+      - tensorflow-io-gcs-filesystem==0.35.0
+      - termcolor==2.4.0
+      - terminado==0.18.0
+      - threadpoolctl==3.2.0
+      - tinycss2==1.2.1
+      - tokenizers==0.15.0
+      - torch==2.4.0.dev20240323+cu121
+      - torchaudio==2.2.0.dev20240323+cu121
+      - torchvision==0.19.0.dev20240323+cu121
+      - tornado==6.4
+      - tqdm==4.66.1
+      - traitlets==5.14.1
+      - transformers==4.39.0.dev0
+      - triton==2.2.0
+      - trl==0.7.7
+      - types-python-dateutil==2.8.19.20240106
+      - typing-extensions==4.10.0
+      - tyro==0.7.0
+      - tzdata==2023.3
+      - uri-template==1.3.0
+      - urllib3==2.1.0
+      - wandb==0.16.1
+      - wcwidth==0.2.13
+      - webcolors==1.13
+      - webencodings==0.5.1
+      - websocket-client==1.7.0
+      - werkzeug==3.0.1
+      - wrapt==1.14.1
+      - xxhash==3.4.1
+      - yarl==1.9.4
+prefix: /fsx/sanchit/miniconda3/envs/venv

wandb/run-20240416_205309-xdytsc71/files/config.yaml

@@ -0,0 +1,40 @@
+wandb_version: 1
+
+_wandb:
+  desc: null
+  value:
+    python_version: 3.11.5
+    cli_version: 0.16.1
+    framework: huggingface
+    huggingface_version: 4.40.0.dev0
+    is_jupyter_run: false
+    is_kaggle_kernel: false
+    start_time: 1713300789.254433
+    t:
+      1:
+      - 1
+      - 11
+      - 49
+      - 51
+      - 55
+      - 71
+      - 98
+      2:
+      - 1
+      - 2
+      - 3
+      - 11
+      - 49
+      - 51
+      - 55
+      - 71
+      - 98
+      3:
+      - 2
+      - 23
+      4: 3.11.5
+      5: 0.16.1
+      6: 4.40.0.dev0
+      8:
+      - 5
+      13: linux-x86_64

wandb/run-20240416_205309-xdytsc71/files/requirements.txt

@@ -0,0 +1,223 @@
+absl-py==2.0.0
+accelerate==0.29.2
+aiohttp==3.9.1
+aiosignal==1.3.1
+annotated-types==0.6.0
+anyio==4.2.0
+appdirs==1.4.4
+argon2-cffi-bindings==21.2.0
+argon2-cffi==23.1.0
+arrow==1.3.0
+asttokens==2.4.1
+astunparse==1.6.3
+async-lru==2.0.4
+attrs==23.1.0
+audioread==3.0.1
+babel==2.14.0
+beautifulsoup4==4.12.3
+bitsandbytes==0.43.1
+bleach==6.1.0
+brotli==1.0.9
+cachetools==5.3.2
+certifi==2023.11.17
+cffi==1.16.0
+chardet==5.2.0
+charset-normalizer==2.0.4
+click==8.1.7
+comm==0.2.1
+cryptography==41.0.7
+datasets==2.18.1.dev0
+debugpy==1.8.1
+decorator==5.1.1
+deepspeed==0.12.2
+defusedxml==0.7.1
+dill==0.3.7
+docker-pycreds==0.4.0
+docstring-parser==0.15
+einops==0.7.0
+evaluate==0.4.0
+executing==2.0.1
+fastjsonschema==2.19.1
+filelock==3.13.1
+flatbuffers==23.5.26
+fqdn==1.5.1
+frozenlist==1.4.1
+fsspec==2023.10.0
+gast==0.5.4
+gitdb==4.0.11
+gitpython==3.1.40
+gmpy2==2.1.2
+google-auth-oauthlib==1.2.0
+google-auth==2.26.1
+google-pasta==0.2.0
+grpcio==1.60.0
+h11==0.14.0
+h5py==3.10.0
+hf-transfer==0.1.5
+hjson==3.1.0
+httpcore==1.0.2
+httpx==0.26.0
+huggingface-hub==0.22.2
+idna==3.4
+ipdb==0.13.13
+ipykernel==6.29.2
+ipython==8.21.0
+isoduration==20.11.0
+jedi==0.19.1
+jinja2==3.1.2
+jiwer==3.0.3
+joblib==1.3.2
+json5==0.9.14
+jsonpointer==2.4
+jsonschema-specifications==2023.12.1
+jsonschema==4.21.1
+jupyter-client==8.6.0
+jupyter-core==5.7.1
+jupyter-events==0.9.0
+jupyter-lsp==2.2.2
+jupyter-server-terminals==0.5.2
+jupyter-server==2.12.5
+jupyterlab-pygments==0.3.0
+jupyterlab-server==2.25.2
+jupyterlab==4.1.1
+keras==2.15.0
+lazy-loader==0.3
+libclang==16.0.6
+librosa==0.10.1
+llvmlite==0.41.1
+markdown-it-py==3.0.0
+markdown==3.5.1
+markupsafe==2.1.1
+matplotlib-inline==0.1.6
+mdurl==0.1.2
+mistune==3.0.2
+mkl-fft==1.3.8
+mkl-random==1.2.4
+mkl-service==2.4.0
+ml-dtypes==0.2.0
+mpmath==1.3.0
+msgpack==1.0.7
+multidict==6.0.4
+multiprocess==0.70.15
+nbclient==0.9.0
+nbconvert==7.16.0
+nbformat==5.9.2
+nest-asyncio==1.6.0
+networkx==3.1
+ninja==1.11.1.1
+nltk==3.8.1
+notebook-shim==0.2.3
+numba==0.58.1
+numpy==1.26.2
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.20.5
+nvidia-nvjitlink-cu12==12.3.101
+nvidia-nvtx-cu12==12.1.105
+oauthlib==3.2.2
+opt-einsum==3.3.0
+overrides==7.7.0
+packaging==23.2
+pandas==2.1.4
+pandocfilters==1.5.1
+parso==0.8.3
+peft==0.7.1
+pexpect==4.9.0
+pillow==10.2.0
+pip==24.0
+platformdirs==4.1.0
+pooch==1.8.0
+prometheus-client==0.19.0
+prompt-toolkit==3.0.43
+protobuf==3.20.2
+psutil==5.9.7
+ptyprocess==0.7.0
+pure-eval==0.2.2
+py-cpuinfo==9.0.0
+pyarrow-hotfix==0.6
+pyarrow==14.0.2
+pyasn1-modules==0.3.0
+pyasn1==0.5.1
+pycparser==2.21
+pydantic-core==2.16.1
+pydantic==2.6.0
+pygments==2.17.2
+pynvml==11.5.0
+pyopenssl==23.2.0
+pysocks==1.7.1
+python-dateutil==2.8.2
+python-json-logger==2.0.7
+pytorch-triton==3.0.0+989adb9a29
+pytz==2023.3.post1
+pyyaml==6.0.1
+pyzmq==25.1.2
+rapidfuzz==3.6.1
+referencing==0.33.0
+regex==2023.12.25
+requests-oauthlib==1.3.1
+requests==2.31.0
+responses==0.18.0
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rich==13.7.0
+rpds-py==0.17.1
+rsa==4.9
+safetensors==0.4.1
+scikit-learn==1.3.2
+scipy==1.11.4
+send2trash==1.8.2
+sentencepiece==0.1.99
+sentry-sdk==1.39.1
+setproctitle==1.3.3
+setuptools==68.2.2
+shtab==1.6.5
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.0
+soundfile==0.12.1
+soupsieve==2.5
+soxr==0.3.7
+stack-data==0.6.3
+sympy==1.12
+tensorboard-data-server==0.7.2
+tensorboard==2.15.1
+tensorflow-cpu==2.15.0.post1
+tensorflow-estimator==2.15.0
+tensorflow-io-gcs-filesystem==0.35.0
+termcolor==2.4.0
+terminado==0.18.0
+threadpoolctl==3.2.0
+tinycss2==1.2.1
+tokenizers==0.15.0
+torch==2.4.0.dev20240323+cu121
+torchaudio==2.2.0.dev20240323+cu121
+torchvision==0.19.0.dev20240323+cu121
+tornado==6.4
+tqdm==4.66.1
+traitlets==5.14.1
+transformers==4.39.0.dev0
+triton==2.2.0
+trl==0.7.7
+types-python-dateutil==2.8.19.20240106
+typing-extensions==4.10.0
+tyro==0.7.0
+tzdata==2023.3
+uri-template==1.3.0
+urllib3==1.26.18
+wandb==0.16.1
+wcwidth==0.2.13
+webcolors==1.13
+webencodings==0.5.1
+websocket-client==1.7.0
+werkzeug==3.0.1
+wheel==0.41.2
+wrapt==1.14.1
+xxhash==3.4.1
+yarl==1.9.4

wandb/run-20240416_205309-xdytsc71/files/wandb-metadata.json

@@ -0,0 +1,552 @@
+{
+    "os": "Linux-5.15.0-1048-aws-x86_64-with-glibc2.31",
+    "python": "3.11.5",
+    "heartbeatAt": "2024-04-16T20:53:09.639270",
+    "startedAt": "2024-04-16T20:53:09.231836",
+    "docker": null,
+    "cuda": null,
+    "args": [
+        "config_mistral.yaml"
+    ],
+    "state": "running",
+    "program": "/fsx/sanchit/mistral-debug-4bit/run_distillation.py",
+    "codePathLocal": "run_distillation.py",
+    "codePath": "run_distillation.py",
+    "host": "ip-26-0-164-187",
+    "username": "sanchit",
+    "executable": "/fsx/sanchit/miniconda3/envs/venv/bin/python",
+    "cpu_count": 96,
+    "cpu_count_logical": 96,
+    "cpu_freq": {
+        "current": 2659.3121145833315,
+        "min": 0.0,
+        "max": 0.0
+    },
+    "cpu_freq_per_core": [
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3544.153,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        }
+    ],
+    "disk": {
+        "/": {
+            "total": 290.7472343444824,
+            "used": 58.411903381347656
+        }
+    },
+    "gpu": "NVIDIA H100 80GB HBM3",
+    "gpu_count": 8,
+    "gpu_devices": [
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        }
+    ],
+    "memory": {
+        "total": 1999.9855461120605
+    }
+}

wandb/run-20240416_205309-xdytsc71/files/wandb-summary.json

@@ -0,0 +1 @@
+{"train/loss": 1.6176064014434814, "train/ce_loss": 1.277328610420227, "train/kl_loss": 0.5957434773445129, "train/time": 67493.65216779709, "train/epoch": 0, "train/learning_rate": 0.0, "_timestamp": 1713368583.0371525, "_runtime": 67793.78271961212, "_step": 50000, "_wandb": {"runtime": 67808}, "cosmopedia/auto_math_text/loss": 1.658003807067871, "cosmopedia/auto_math_text/ce_loss": 1.3764740228652954, "cosmopedia/auto_math_text/kl_loss": 0.5568245649337769, "cosmopedia/auto_math_text/perplexity": 3.960910895239512, "cosmopedia/auto_math_text/time": 13.17593264579773, "cosmopedia/auto_math_text/epoch": 0, "cosmopedia/khanacademy/loss": 1.413199782371521, "cosmopedia/khanacademy/ce_loss": 1.0476577281951904, "cosmopedia/khanacademy/kl_loss": 0.5750735998153687, "cosmopedia/khanacademy/perplexity": 2.85096555518061, "cosmopedia/khanacademy/time": 13.071500062942505, "cosmopedia/khanacademy/epoch": 0, "cosmopedia/openstax/loss": 1.6419086456298828, "cosmopedia/openstax/ce_loss": 1.3129570484161377, "cosmopedia/openstax/kl_loss": 0.5915430188179016, "cosmopedia/openstax/perplexity": 3.717149266941545, "cosmopedia/openstax/time": 12.288958072662354, "cosmopedia/openstax/epoch": 0, "cosmopedia/stanford/loss": 1.461927890777588, "cosmopedia/stanford/ce_loss": 1.1311124563217163, "cosmopedia/stanford/kl_loss": 0.5570380091667175, "cosmopedia/stanford/perplexity": 3.099102199007535, "cosmopedia/stanford/time": 12.182331085205078, "cosmopedia/stanford/epoch": 0, "cosmopedia/stories/loss": 1.6099693775177002, "cosmopedia/stories/ce_loss": 1.308800458908081, "cosmopedia/stories/kl_loss": 0.5629290342330933, "cosmopedia/stories/perplexity": 3.70173066988729, "cosmopedia/stories/time": 13.002483367919922, "cosmopedia/stories/epoch": 0, "cosmopedia/web_samples_v1/loss": 1.58097243309021, "cosmopedia/web_samples_v1/ce_loss": 1.3564168214797974, "cosmopedia/web_samples_v1/kl_loss": 0.49583905935287476, "cosmopedia/web_samples_v1/perplexity": 3.882257527739104, "cosmopedia/web_samples_v1/time": 12.979966402053833, "cosmopedia/web_samples_v1/epoch": 0, "cosmopedia/web_samples_v2/loss": 1.638546109199524, "cosmopedia/web_samples_v2/ce_loss": 1.3901900053024292, "cosmopedia/web_samples_v2/kl_loss": 0.5263941287994385, "cosmopedia/web_samples_v2/perplexity": 4.015612968269408, "cosmopedia/web_samples_v2/time": 13.090406656265259, "cosmopedia/web_samples_v2/epoch": 0, "cosmopedia/wikihow/loss": 1.490721344947815, "cosmopedia/wikihow/ce_loss": 1.2553372383117676, "cosmopedia/wikihow/kl_loss": 0.48645156621932983, "cosmopedia/wikihow/perplexity": 3.5090215513289906, "cosmopedia/wikihow/time": 13.000000238418579, "cosmopedia/wikihow/epoch": 0}

wandb/run-20240416_205309-xdytsc71/logs/debug-internal.log

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

wandb/run-20240416_205309-xdytsc71/logs/debug.log

@@ -0,0 +1,34 @@
+2024-04-16 20:53:09,244 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Current SDK version is 0.16.1
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Configure stats pid to 965482
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Loading settings from /admin/home/sanchit/.config/wandb/settings
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Loading settings from /fsx/sanchit/mistral-debug-4bit/wandb/settings
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Loading settings from environment variables: {}
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Applying setup settings: {'_disable_service': False}
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_setup.py:_flush():76] Inferring run settings from compute environment: {'program_relpath': 'run_distillation.py', 'program_abspath': '/fsx/sanchit/mistral-debug-4bit/run_distillation.py', 'program': '/fsx/sanchit/mistral-debug-4bit/run_distillation.py'}
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:_log_setup():524] Logging user logs to /fsx/sanchit/mistral-debug-4bit/wandb/run-20240416_205309-xdytsc71/logs/debug.log
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:_log_setup():525] Logging internal logs to /fsx/sanchit/mistral-debug-4bit/wandb/run-20240416_205309-xdytsc71/logs/debug-internal.log
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:init():564] calling init triggers
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:init():571] wandb.init called with sweep_config: {}
+config: {}
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:init():614] starting backend
+2024-04-16 20:53:09,245 INFO    MainThread:965482 [wandb_init.py:init():618] setting up manager
+2024-04-16 20:53:09,248 INFO    MainThread:965482 [backend.py:_multiprocessing_setup():105] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-04-16 20:53:09,254 INFO    MainThread:965482 [wandb_init.py:init():624] backend started and connected
+2024-04-16 20:53:09,257 INFO    MainThread:965482 [wandb_init.py:init():716] updated telemetry
+2024-04-16 20:53:09,257 INFO    MainThread:965482 [wandb_init.py:init():749] communicating run to backend with 90.0 second timeout
+2024-04-16 20:53:09,461 INFO    MainThread:965482 [wandb_run.py:_on_init():2254] communicating current version
+2024-04-16 20:53:09,520 INFO    MainThread:965482 [wandb_run.py:_on_init():2263] got version response upgrade_message: "wandb version 0.16.6 is available!  To upgrade, please run:\n $ pip install wandb --upgrade"
+
+2024-04-16 20:53:09,520 INFO    MainThread:965482 [wandb_init.py:init():800] starting run threads in backend
+2024-04-16 20:53:15,732 INFO    MainThread:965482 [wandb_run.py:_console_start():2233] atexit reg
+2024-04-16 20:53:15,732 INFO    MainThread:965482 [wandb_run.py:_redirect():2088] redirect: wrap_raw
+2024-04-16 20:53:15,732 INFO    MainThread:965482 [wandb_run.py:_redirect():2153] Wrapping output streams.
+2024-04-16 20:53:15,732 INFO    MainThread:965482 [wandb_run.py:_redirect():2178] Redirects installed.
+2024-04-16 20:53:15,733 INFO    MainThread:965482 [wandb_init.py:init():841] run started, returning control to user process
+2024-04-17 15:43:17,904 INFO    MainThread:965482 [wandb_run.py:_finish():1962] finishing run sanchit-gandhi/distil-mixtral/xdytsc71
+2024-04-17 15:43:17,906 INFO    MainThread:965482 [wandb_run.py:_atexit_cleanup():2202] got exitcode: 0
+2024-04-17 15:43:17,906 INFO    MainThread:965482 [wandb_run.py:_restore():2185] restore
+2024-04-17 15:43:17,907 INFO    MainThread:965482 [wandb_run.py:_restore():2191] restore done
+2024-04-17 15:43:24,105 INFO    MainThread:965482 [wandb_run.py:_footer_history_summary_info():3837] rendering history
+2024-04-17 15:43:24,106 INFO    MainThread:965482 [wandb_run.py:_footer_history_summary_info():3869] rendering summary
+2024-04-17 15:43:24,118 INFO    MainThread:965482 [wandb_run.py:_footer_sync_info():3796] logging synced files

wandb/run-20240416_205309-xdytsc71/run-xdytsc71.wandb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cec0ae8a1202accebb26f3be48ab20552867b6b02ecd7ec9bb685c2b6131603
+size 26549109