Upload 2 files

app.py

@@ -0,0 +1,367 @@
+"""
+This module is used to launch Axolotl with user defined configurations.
+"""
+
+import gradio as gr
+import yaml
+from config import config
+
+example_yml = """
+base_model: NousResearch/Llama-2-7b-hf
+model_type: LlamaForCausalLM
+tokenizer_type: LlamaTokenizer
+
+load_in_8bit: false
+load_in_4bit: true
+strict: false
+
+datasets:
+  - path: mhenrichsen/alpaca_2k_test
+    type: alpaca
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./qlora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 4096
+sample_packing: true
+pad_to_sequence_len: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules:
+lora_target_linear: true
+lora_fan_in_fan_out:
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 4
+optimizer: paged_adamw_32bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+train_on_inputs: false
+group_by_length: false
+bf16: auto
+fp16:
+tf32: false
+
+gradient_checkpointing: true
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:
+logging_steps: 1
+xformers_attention:
+flash_attention: true
+
+warmup_steps: 10
+evals_per_epoch: 4
+eval_table_size:
+saves_per_epoch: 1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:
+special_tokens:
+"""
+
+
+def yml_config(yml_config):
+  """
+    This function saves as a yaml file from user text.
+    """
+  yml_config = yaml.safe_load(yml_config)
+  with open("config.yml", "w", encoding="utf-8") as file:
+    yaml.dump(yml_config, file)
+  # print(yml_config)
+  return yaml.dump(yml_config)
+
+
+with gr.Blocks(title="Axolotl Launcher") as demo:
+  gr.Markdown("""
+    # Axolotl Launcher
+    Fill out the required fields below to create a training run.
+    """)
+  with gr.Tab("Base Model & Tokenizer"):
+    with gr.Column():
+      with gr.Row():
+        base_model = gr.Textbox(label="Base Model")
+        base_model_ignore_patterns = gr.Textbox(
+            label="Base Model Ignore Patterns")
+        base_model_config = gr.Textbox(label="Base Model Config")
+        model_revision = gr.Textbox(label="Model Revision")
+      with gr.Row():
+        tokenizer_config = gr.Textbox(label="Tokenizer Config")
+        model_type = gr.Textbox(label="Model Type")
+        tokenizer_type = gr.Textbox(label="Tokenizer Type")
+      with gr.Row():
+        trust_remote_code = gr.Checkbox(label="Trust Remote Code", value=False)
+        tokenizer_use_fast = gr.Checkbox(label="Use Fast Tokenizer",
+                                         value=True)
+        tokenizer_legacy = gr.Checkbox(label="Use Legacy Tokenizer",
+                                       value=False)
+        resize_token_embeddings_to_32x = gr.Checkbox(
+            label="Resize Token Embeddings to 32x", value=False)
+    with gr.Accordion("Adv. Config", open=False):
+      with gr.Tab("Model Derivation & Configuration Overrides"):
+        with gr.Column():
+          is_falcon_derived_model = gr.Checkbox(
+              label="Is Falcon Derived Model", value=False)
+          is_llama_derived_model = gr.Checkbox(label="Is Llama Derived Model",
+                                               value=False)
+          is_mistral_derived_model = gr.Checkbox(
+              label="Is Mistral Derived Model", value=False)
+          is_qwen_derived_model = gr.Checkbox(label="Is Qwen Derived Model",
+                                              value=False)
+          model_config = gr.TextArea(label="Model Config Overrides",
+                                     placeholder="YAML or JSON format")
+          bnb_config_kwargs = gr.TextArea(label="BnB Config KWArgs",
+                                          placeholder="YAML or JSON format")
+
+      with gr.Tab("Quantization & Precision"):
+        with gr.Column():
+          with gr.Row():
+            gptq = gr.Checkbox(label="GPTQ", value=False)
+            gptq_groupsize = gr.Number(label="GPTQ Groupsize", value=128)
+            gptq_model_v1 = gr.Checkbox(label="GPTQ Model V1", value=False)
+            load_in_8bit = gr.Checkbox(label="Load in 8-bit", value=False)
+            load_in_4bit = gr.Checkbox(label="Load in 4-bit", value=False)
+          with gr.Row():
+            bf16 = gr.Checkbox(label="BF16", value=False)
+            fp16 = gr.Checkbox(label="FP16", value=False)
+            tf32 = gr.Checkbox(label="TF32", value=False)
+            bfloat16 = gr.Checkbox(label="BFloat16", value=False)
+            float16 = gr.Checkbox(label="Float16", value=False)
+
+      with gr.Tab("GPU & LoRA Settings"):
+        gpu_memory_limit = gr.Textbox(label="GPU Memory Limit")
+        lora_on_cpu = gr.Checkbox(label="LoRA on CPU", value=False)
+        datasets = gr.TextArea(label="Datasets",
+                               placeholder="YAML or JSON format for datasets")
+        test_datasets = gr.TextArea(
+            label="Test Datasets",
+            placeholder="YAML or JSON format for test datasets")
+        rl = gr.Textbox(label="RL")
+        chat_template = gr.Textbox(label="Chat Template")
+        default_system_message = gr.Textbox(label="Default System Message")
+        dataset_prepared_path = gr.Textbox(label="Dataset Prepared Path")
+        push_dataset_to_hub = gr.Textbox(label="Push Dataset to Hub")
+        dataset_processes = gr.Number(label="Dataset Processes", value=1)
+        dataset_keep_in_memory = gr.Checkbox(label="Dataset Keep in Memory",
+                                             value=False)
+        with gr.Row():
+          hub_model_id = gr.Textbox(label="Hub Model ID")
+          hub_strategy = gr.Textbox(label="Hub Strategy")
+          hf_use_auth_token = gr.Checkbox(label="HF Use Auth Token",
+                                          value=False)
+        with gr.Row():
+          val_set_size = gr.Number(label="Validation Set Size",
+                                   value=0.04,
+                                   step=0.01)
+          dataset_shard_num = gr.Number(label="Dataset Shard Num")
+          dataset_shard_idx = gr.Number(label="Dataset Shard Index")
+
+      with gr.Tab("Training & Evaluation"):
+        with gr.Row():
+          sequence_len = gr.Number(label="Sequence Length", value=2048)
+          pad_to_sequence_len = gr.Checkbox(label="Pad to Sequence Length",
+                                            value=False)
+        with gr.Row():
+          sample_packing = gr.Checkbox(label="Sample Packing", value=False)
+          eval_sample_packing = gr.Checkbox(label="Eval Sample Packing",
+                                            value=False)
+          sample_packing_eff_est = gr.Number(label="Sample Packing Eff Est")
+        with gr.Row():
+          total_num_tokens = gr.Number(label="Total Num Tokens")
+          device_map = gr.Textbox(label="Device Map")
+          max_memory = gr.Textbox(label="Max Memory")
+          adapter = gr.Textbox(label="Adapter")
+        with gr.Column():
+          lora_model_dir = gr.Textbox(label="LoRA Model Dir")
+          lora_r = gr.Number(label="LoRA R", value=8)
+          lora_alpha = gr.Number(label="LoRA Alpha", value=16)
+          lora_dropout = gr.Number(label="LoRA Dropout", value=0.05, step=0.01)
+          lora_target_modules = gr.TextArea(label="LoRA Target Modules")
+          lora_target_linear = gr.Checkbox(label="LoRA Target Linear",
+                                           value=False)
+          lora_modules_to_save = gr.TextArea(label="LoRA Modules to Save")
+          lora_fan_in_fan_out = gr.Checkbox(label="LoRA Fan In Fan Out",
+                                            value=False)
+          peft = gr.Textbox(label="PEFT")
+        with gr.Row():
+          relora_steps = gr.Number(label="ReLoRA Steps")
+          relora_warmup_steps = gr.Number(label="ReLoRA Warmup Steps")
+          relora_anneal_steps = gr.Number(label="ReLoRA Anneal Steps")
+          relora_prune_ratio = gr.Number(label="ReLoRA Prune Ratio")
+          relora_cpu_offload = gr.Checkbox(label="ReLoRA CPU Offload",
+                                           value=False)
+        with gr.Row():
+          wandb_mode = gr.Textbox(label="WandB Mode")
+          wandb_project = gr.Textbox(label="WandB Project")
+          wandb_entity = gr.Textbox(label="WandB Entity")
+          wandb_watch = gr.Checkbox(label="WandB Watch", value=False)
+          wandb_name = gr.Textbox(label="WandB Name")
+          wandb_run_id = gr.Textbox(label="WandB Run ID")
+          wandb_log_model = gr.Checkbox(label="WandB Log Model", value=False)
+        with gr.Column():
+          mlflow_tracking_uri = gr.Textbox(label="MLFlow Tracking URI")
+          mlflow_experiment_name = gr.Textbox(label="MLFlow Experiment Name")
+          output_dir = gr.Textbox(label="Output Dir")
+          torch_compile = gr.Checkbox(label="Torch Compile", value=False)
+          torch_compile_backend = gr.Textbox(label="Torch Compile Backend")
+          gradient_accumulation_steps = gr.Number(
+              label="Gradient Accumulation Steps", value=1)
+          micro_batch_size = gr.Number(label="Micro Batch Size", value=2)
+          eval_batch_size = gr.Number(label="Eval Batch Size", value=2)
+          num_epochs = gr.Number(label="Number of Epochs", value=4)
+          warmup_steps = gr.Number(label="Warmup Steps", value=100)
+          warmup_ratio = gr.Number(label="Warmup Ratio")
+          learning_rate = gr.Number(label="Learning Rate",
+                                    value=0.00003,
+                                    step=1e-5)
+          lr_quadratic_warmup = gr.Checkbox(label="LR Quadratic Warmup",
+                                            value=False)
+          logging_steps = gr.Number(label="Logging Steps", value=1)
+          eval_steps = gr.Textbox(label="Eval Steps")
+          evals_per_epoch = gr.Number(label="Evals per Epoch", value=4)
+          save_strategy = gr.Textbox(label="Save Strategy")
+          save_steps = gr.Textbox(label="Save Steps")
+          saves_per_epoch = gr.Number(label="Saves per Epoch", value=1)
+          save_total_limit = gr.Number(label="Save Total Limit")
+          max_steps = gr.Number(label="Max Steps")
+          eval_table_size = gr.Number(label="Eval Table Size")
+          eval_max_new_tokens = gr.Number(label="Eval Max New Tokens",
+                                          value=128)
+          eval_causal_lm_metrics = gr.TextArea(label="Eval Causal LM Metrics")
+          loss_watchdog_threshold = gr.Number(label="Loss Watchdog Threshold")
+          loss_watchdog_patience = gr.Number(label="Loss Watchdog Patience",
+                                             value=3)
+          save_safetensors = gr.Checkbox(label="Save SafeTensors", value=False)
+          train_on_inputs = gr.Checkbox(label="Train on Inputs", value=False)
+          group_by_length = gr.Checkbox(label="Group by Length", value=False)
+          gradient_checkpointing = gr.Checkbox(label="Gradient Checkpointing",
+                                               value=False)
+          early_stopping_patience = gr.Number(label="Early Stopping Patience",
+                                              value=3)
+          lr_scheduler = gr.Textbox(label="LR Scheduler")
+          lr_scheduler_kwargs = gr.TextArea(label="LR Scheduler KWArgs")
+          cosine_min_lr_ratio = gr.Number(label="Cosine Min LR Ratio")
+          cosine_constant_lr_ratio = gr.Number(
+              label="Cosine Constant LR Ratio")
+          lr_div_factor = gr.Number(label="LR Div Factor")
+          log_sweep_min_lr = gr.Number(label="Log Sweep Min LR")
+          log_sweep_max_lr = gr.Number(label="Log Sweep Max LR")
+          optimizer = gr.Textbox(label="Optimizer")
+          weight_decay = gr.Number(label="Weight Decay", value=0.0, step=0.01)
+          adam_beta1 = gr.Number(label="Adam Beta1", value=0.9, step=0.01)
+          adam_beta2 = gr.Number(label="Adam Beta2", value=0.999, step=0.001)
+          adam_epsilon = gr.Number(label="Adam Epsilon", value=1e-8, step=1e-9)
+          max_grad_norm = gr.Number(label="Max Grad Norm")
+          neftune_noise_alpha = gr.Number(label="NEFTune Noise Alpha")
+          flash_optimum = gr.Checkbox(label="Flash Optimum", value=False)
+          xformers_attention = gr.Checkbox(label="XFormers Attention",
+                                           value=False)
+          flash_attention = gr.Checkbox(label="Flash Attention", value=False)
+          flash_attn_cross_entropy = gr.Checkbox(
+              label="Flash Attn Cross Entropy", value=False)
+          flash_attn_rms_norm = gr.Checkbox(label="Flash Attn RMS Norm",
+                                            value=False)
+          flash_attn_fuse_qkv = gr.Checkbox(label="Flash Attn Fuse QKV",
+                                            value=False)
+          flash_attn_fuse_mlp = gr.Checkbox(label="Flash Attn Fuse MLP",
+                                            value=False)
+          sdp_attention = gr.Checkbox(label="SDP Attention", value=False)
+          s2_attention = gr.Checkbox(label="S2 Attention", value=False)
+          resume_from_checkpoint = gr.Textbox(label="Resume From Checkpoint")
+          auto_resume_from_checkpoints = gr.Checkbox(
+              label="Auto Resume From Checkpoints", value=False)
+          local_rank = gr.Number(label="Local Rank")
+          special_tokens = gr.TextArea(label="Special Tokens")
+          tokens = gr.TextArea(label="Tokens")
+          fsdp = gr.Checkbox(label="FSDP", value=False)
+          fsdp_config = gr.TextArea(label="FSDP Config")
+          deepspeed = gr.Textbox(label="Deepspeed")
+          ddp_timeout = gr.Number(label="DDP Timeout")
+          ddp_bucket_cap_mb = gr.Number(label="DDP Bucket Cap MB")
+          ddp_broadcast_buffers = gr.Checkbox(label="DDP Broadcast Buffers",
+                                              value=False)
+          torchdistx_path = gr.Textbox(label="TorchDistX Path")
+          pretraining_dataset = gr.Textbox(label="Pretraining Dataset")
+          debug = gr.Checkbox(label="Debug", value=False)
+          seed = gr.Number(label="Seed", value=42)
+          strict = gr.Checkbox(label="Strict", value=False)
+
+    submit_button = gr.Button("Launch Configuration")
+    output_area = gr.TextArea(label="Configuration Output")
+
+  submit_button.click(
+      config,
+      inputs=[
+          base_model, base_model_ignore_patterns, base_model_config,
+          model_revision, tokenizer_config, model_type, tokenizer_type,
+          trust_remote_code, tokenizer_use_fast, tokenizer_legacy,
+          resize_token_embeddings_to_32x, is_falcon_derived_model,
+          is_llama_derived_model, is_mistral_derived_model,
+          is_qwen_derived_model, model_config, bnb_config_kwargs, gptq,
+          gptq_groupsize, gptq_model_v1, load_in_8bit, load_in_4bit, bf16,
+          fp16, tf32, bfloat16, float16, gpu_memory_limit, lora_on_cpu,
+          datasets, test_datasets, rl, chat_template, default_system_message,
+          dataset_prepared_path, push_dataset_to_hub, dataset_processes,
+          dataset_keep_in_memory, hub_model_id, hub_strategy,
+          hf_use_auth_token, val_set_size, dataset_shard_num,
+          dataset_shard_idx, sequence_len, pad_to_sequence_len, sample_packing,
+          eval_sample_packing, sample_packing_eff_est, total_num_tokens,
+          device_map, max_memory, adapter, lora_model_dir, lora_r, lora_alpha,
+          lora_dropout, lora_target_modules, lora_target_linear,
+          lora_modules_to_save, lora_fan_in_fan_out, peft, relora_steps,
+          relora_warmup_steps, relora_anneal_steps, relora_prune_ratio,
+          relora_cpu_offload, wandb_mode, wandb_project, wandb_entity,
+          wandb_watch, wandb_name, wandb_run_id, wandb_log_model,
+          mlflow_tracking_uri, mlflow_experiment_name, output_dir,
+          torch_compile, torch_compile_backend, gradient_accumulation_steps,
+          micro_batch_size, eval_batch_size, num_epochs, warmup_steps,
+          warmup_ratio, learning_rate, lr_quadratic_warmup, logging_steps,
+          eval_steps, evals_per_epoch, save_strategy, save_steps,
+          saves_per_epoch, save_total_limit, max_steps, eval_table_size,
+          eval_max_new_tokens, eval_causal_lm_metrics, loss_watchdog_threshold,
+          loss_watchdog_patience, save_safetensors, train_on_inputs,
+          group_by_length, gradient_checkpointing, early_stopping_patience,
+          lr_scheduler, lr_scheduler_kwargs, cosine_min_lr_ratio,
+          cosine_constant_lr_ratio, lr_div_factor, log_sweep_min_lr,
+          log_sweep_max_lr, optimizer, weight_decay, adam_beta1, adam_beta2,
+          adam_epsilon, max_grad_norm, neftune_noise_alpha, flash_optimum,
+          xformers_attention, flash_attention, flash_attn_cross_entropy,
+          flash_attn_rms_norm, flash_attn_fuse_qkv, flash_attn_fuse_mlp,
+          sdp_attention, s2_attention, resume_from_checkpoint,
+          auto_resume_from_checkpoints, local_rank, special_tokens, tokens,
+          fsdp, fsdp_config, deepspeed, ddp_timeout, ddp_bucket_cap_mb,
+          ddp_broadcast_buffers, torchdistx_path, pretraining_dataset, debug,
+          seed, strict
+      ],
+      outputs=output_area)
+  """
+    This section is used to create a configuration file from user text.
+  """
+  with gr.Tab(label="YML text"):
+    yml_config_text = gr.TextArea(label='YML Config',
+                                  lines=50,
+                                  value=example_yml)
+    create_config = gr.Button("Create config")
+    output = gr.TextArea(label="Generated config")
+    create_config.click(
+        yml_config,
+        inputs=[yml_config_text],
+        outputs=output,
+    )
+
+demo.launch(share=True)

config.py

@@ -0,0 +1,200 @@
+import yaml
+
+
+def config(
+    base_model, base_model_ignore_patterns, base_model_config, model_revision,
+    tokenizer_config, model_type, tokenizer_type, trust_remote_code,
+    tokenizer_use_fast, tokenizer_legacy, resize_token_embeddings_to_32x,
+    is_falcon_derived_model, is_llama_derived_model, is_mistral_derived_model,
+    is_qwen_derived_model, model_config, bnb_config_kwargs, gptq,
+    gptq_groupsize, gptq_model_v1, load_in_8bit, load_in_4bit, bf16, fp16,
+    tf32, bfloat16, float16, gpu_memory_limit, lora_on_cpu, datasets,
+    test_datasets, rl, chat_template, default_system_message,
+    dataset_prepared_path, push_dataset_to_hub, dataset_processes,
+    dataset_keep_in_memory, hub_model_id, hub_strategy, hf_use_auth_token,
+    val_set_size, dataset_shard_num, dataset_shard_idx, sequence_len,
+    pad_to_sequence_len, sample_packing, eval_sample_packing,
+    sample_packing_eff_est, total_num_tokens, device_map, max_memory, adapter,
+    lora_model_dir, lora_r, lora_alpha, lora_dropout, lora_target_modules,
+    lora_target_linear, lora_modules_to_save, lora_fan_in_fan_out, peft,
+    relora_steps, relora_warmup_steps, relora_anneal_steps, relora_prune_ratio,
+    relora_cpu_offload, wandb_mode, wandb_project, wandb_entity, wandb_watch,
+    wandb_name, wandb_run_id, wandb_log_model, mlflow_tracking_uri,
+    mlflow_experiment_name, output_dir, torch_compile, torch_compile_backend,
+    gradient_accumulation_steps, micro_batch_size, eval_batch_size, num_epochs,
+    warmup_steps, warmup_ratio, learning_rate, lr_quadratic_warmup,
+    logging_steps, eval_steps, evals_per_epoch, save_strategy, save_steps,
+    saves_per_epoch, save_total_limit, max_steps, eval_table_size,
+    eval_max_new_tokens, eval_causal_lm_metrics, loss_watchdog_threshold,
+    loss_watchdog_patience, save_safetensors, train_on_inputs, group_by_length,
+    gradient_checkpointing, early_stopping_patience, lr_scheduler,
+    lr_scheduler_kwargs, cosine_min_lr_ratio, cosine_constant_lr_ratio,
+    lr_div_factor, log_sweep_min_lr, log_sweep_max_lr, optimizer, weight_decay,
+    adam_beta1, adam_beta2, adam_epsilon, max_grad_norm, neftune_noise_alpha,
+    flash_optimum, xformers_attention, flash_attention,
+    flash_attn_cross_entropy, flash_attn_rms_norm, flash_attn_fuse_qkv,
+    flash_attn_fuse_mlp, sdp_attention, s2_attention, resume_from_checkpoint,
+    auto_resume_from_checkpoints, local_rank, special_tokens, tokens, fsdp,
+    fsdp_config, deepspeed, ddp_timeout, ddp_bucket_cap_mb,
+    ddp_broadcast_buffers, torchdistx_path, pretraining_dataset, debug, seed,
+    strict):
+  """
+      This function generates a configuration dictionary based on the provided parameters and saves it as a yaml file.
+      """
+  config_dict = {
+      # Base model configurations
+      "base_model": base_model,
+      "base_model_ignore_patterns": base_model_ignore_patterns,
+      "base_model_config": base_model_config,
+      "model_revision": model_revision,
+      "tokenizer_config": tokenizer_config,
+      "model_type": model_type,
+      "tokenizer_type": tokenizer_type,
+      "trust_remote_code": trust_remote_code,
+      "tokenizer_use_fast": tokenizer_use_fast,
+      "tokenizer_legacy": tokenizer_legacy,
+      "resize_token_embeddings_to_32x": resize_token_embeddings_to_32x,
+      # Derived model flags
+      "is_falcon_derived_model": is_falcon_derived_model,
+      "is_llama_derived_model": is_llama_derived_model,
+      "is_mistral_derived_model": is_mistral_derived_model,
+      "is_qwen_derived_model": is_qwen_derived_model,
+      # Model configuration overrides
+      "model_config": model_config,
+      "bnb_config_kwargs": bnb_config_kwargs,
+      # Quantization and precision settings
+      "gptq": gptq,
+      "gptq_groupsize": gptq_groupsize,
+      "gptq_model_v1": gptq_model_v1,
+      "load_in_8bit": load_in_8bit,
+      "load_in_4bit": load_in_4bit,
+      "bf16": bf16,
+      "fp16": fp16,
+      "tf32": tf32,
+      "bfloat16": bfloat16,
+      "float16": float16,
+      "gpu_memory_limit": gpu_memory_limit,
+      "lora_on_cpu": lora_on_cpu,
+      # Dataset configurations
+      "datasets": datasets,
+      "test_datasets": test_datasets,
+      "rl": rl,
+      "chat_template": chat_template,
+      "default_system_message": default_system_message,
+      "dataset_prepared_path": dataset_prepared_path,
+      "push_dataset_to_hub": push_dataset_to_hub,
+      "dataset_processes": dataset_processes,
+      "dataset_keep_in_memory": dataset_keep_in_memory,
+      "hub_model_id": hub_model_id,
+      "hub_strategy": hub_strategy,
+      "hf_use_auth_token": hf_use_auth_token,
+      "val_set_size": val_set_size,
+      "dataset_shard_num": dataset_shard_num,
+      "dataset_shard_idx": dataset_shard_idx,
+      # Training hyperparameters
+      "sequence_len": sequence_len,
+      "pad_to_sequence_len": pad_to_sequence_len,
+      "sample_packing": sample_packing,
+      "eval_sample_packing": eval_sample_packing,
+      "sample_packing_eff_est": sample_packing_eff_est,
+      "total_num_tokens": total_num_tokens,
+      "device_map": device_map,
+      "max_memory": max_memory,
+      # Adapter and LoRA settings
+      "adapter": adapter,
+      "lora_model_dir": lora_model_dir,
+      "lora_r": lora_r,
+      "lora_alpha": lora_alpha,
+      "lora_dropout": lora_dropout,
+      "lora_target_modules": lora_target_modules,
+      "lora_target_linear": lora_target_linear,
+      "lora_modules_to_save": lora_modules_to_save,
+      "lora_fan_in_fan_out": lora_fan_in_fan_out,
+      "peft": peft,
+      "relora_steps": relora_steps,
+      "relora_warmup_steps": relora_warmup_steps,
+      "relora_anneal_steps": relora_anneal_steps,
+      "relora_prune_ratio": relora_prune_ratio,
+      "relora_cpu_offload": relora_cpu_offload,
+      # wandb and mlflow configurations
+      "wandb_mode": wandb_mode,
+      "wandb_project": wandb_project,
+      "wandb_entity": wandb_entity,
+      "wandb_watch": wandb_watch,
+      "wandb_name": wandb_name,
+      "wandb_run_id": wandb_run_id,
+      "wandb_log_model": wandb_log_model,
+      "mlflow_tracking_uri": mlflow_tracking_uri,
+      "mlflow_experiment_name": mlflow_experiment_name,
+      "output_dir": output_dir,
+      "torch_compile": torch_compile,
+      "torch_compile_backend": torch_compile_backend,
+      "gradient_accumulation_steps": gradient_accumulation_steps,
+      "micro_batch_size": micro_batch_size,
+      "eval_batch_size": eval_batch_size,
+      "num_epochs": num_epochs,
+      "warmup_steps": warmup_steps,
+      "warmup_ratio": warmup_ratio,
+      "learning_rate": learning_rate,
+      "lr_quadratic_warmup": lr_quadratic_warmup,
+      "logging_steps": logging_steps,
+      "eval_steps": eval_steps,
+      "evals_per_epoch": evals_per_epoch,
+      "save_strategy": save_strategy,
+      "save_steps": save_steps,
+      "saves_per_epoch": saves_per_epoch,
+      "save_total_limit": save_total_limit,
+      "max_steps": max_steps,
+      "eval_table_size": eval_table_size,
+      "eval_max_new_tokens": eval_max_new_tokens,
+      "eval_causal_lm_metrics": eval_causal_lm_metrics,
+      "loss_watchdog_threshold": loss_watchdog_threshold,
+      "loss_watchdog_patience": loss_watchdog_patience,
+      "save_safetensors": save_safetensors,
+      "train_on_inputs": train_on_inputs,
+      "group_by_length": group_by_length,
+      "gradient_checkpointing": gradient_checkpointing,
+      "early_stopping_patience": early_stopping_patience,
+      "lr_scheduler": lr_scheduler,
+      "lr_scheduler_kwargs": lr_scheduler_kwargs,
+      "cosine_min_lr_ratio": cosine_min_lr_ratio,
+      "cosine_constant_lr_ratio": cosine_constant_lr_ratio,
+      "lr_div_factor": lr_div_factor,
+      "log_sweep_min_lr": log_sweep_min_lr,
+      "log_sweep_max_lr": log_sweep_max_lr,
+      "optimizer": optimizer,
+      "weight_decay": weight_decay,
+      "adam_beta1": adam_beta1,
+      "adam_beta2": adam_beta2,
+      "adam_epsilon": adam_epsilon,
+      "max_grad_norm": max_grad_norm,
+      "neftune_noise_alpha": neftune_noise_alpha,
+      "flash_optimum": flash_optimum,
+      "xformers_attention": xformers_attention,
+      "flash_attention": flash_attention,
+      "flash_attn_cross_entropy": flash_attn_cross_entropy,
+      "flash_attn_rms_norm": flash_attn_rms_norm,
+      "flash_attn_fuse_qkv": flash_attn_fuse_qkv,
+      "flash_attn_fuse_mlp": flash_attn_fuse_mlp,
+      "sdp_attention": sdp_attention,
+      "s2_attention": s2_attention,
+      "resume_from_checkpoint": resume_from_checkpoint,
+      "auto_resume_from_checkpoints": auto_resume_from_checkpoints,
+      "local_rank": local_rank,
+      "special_tokens": special_tokens,
+      "tokens": tokens,
+      "fsdp": fsdp,
+      "fsdp_config": fsdp_config,
+      "deepspeed": deepspeed,
+      "ddp_timeout": ddp_timeout,
+      "ddp_bucket_cap_mb": ddp_bucket_cap_mb,
+      "ddp_broadcast_buffers": ddp_broadcast_buffers,
+      "torchdistx_path": torchdistx_path,
+      "pretraining_dataset": pretraining_dataset,
+      "debug": debug,
+      "seed": seed,
+      "strict": strict,
+  }
+  with open("config.yml", "w", encoding="utf-8") as file:
+    yaml.dump(config_dict, file)
+  return yaml.dump(config_dict)