fusion-bench 0.2.6__py3-none-any.whl → 0.2.7__py3-none-any.whl

fusion_bench/utils/plot/token.py

@@ -0,0 +1,52 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+
+
+def visualize_model_inputs(input_ids, attention_mask, labels, tokenizer=None):
+    """
+    Visualize model inputs: attention mask, labels and input_ids
+
+    Parameters:
+    -----------
+    attention_mask: numpy array or tensor
+        The attention mask array
+    labels: numpy array or tensor
+        The labels array
+    input_ids: numpy array or tensor
+        The input ids array
+    tokenizer: optional
+        The tokenizer object to decode input_ids
+    """
+
+    # Convert inputs to numpy if they're tensors
+    attention_mask = np.array(attention_mask)
+    labels = np.array(labels)
+    input_ids = np.array(input_ids)
+
+    # Create figure with 3 subplots
+    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(15, 10))
+
+    # Plot attention mask
+    sns.heatmap(attention_mask.reshape(1, -1), ax=ax1, cmap="Blues", cbar=True)
+    ax1.set_title("**Attention Mask**")
+    ax1.set_ylabel("Sequence")
+
+    # Plot labels
+    sns.heatmap(labels.reshape(1, -1), ax=ax2, cmap="Reds", cbar=True)
+    ax2.set_title("**Labels**")
+    ax2.set_ylabel("Sequence")
+
+    # Plot input_ids
+    sns.heatmap(input_ids.reshape(1, -1), ax=ax3, cmap="Greens", cbar=True)
+    ax3.set_title("**Input IDs**")
+    ax3.set_ylabel("Sequence")
+
+    # If tokenizer is provided, add decoded tokens as x-axis labels
+    if tokenizer:
+        decoded_tokens = [tokenizer.decode(token_id) for token_id in input_ids]
+        ax3.set_xticks(np.arange(len(decoded_tokens)) + 0.5)
+        ax3.set_xticklabels(decoded_tokens, rotation=45, ha="right")
+
+    plt.tight_layout()
+    return fig

fusion_bench/utils/plot/token_notebook.py

@@ -0,0 +1,127 @@
+import numpy as np
+from IPython.display import HTML, display
+
+
+def create_color_style():
+    return """
+    <style>
+        .token-container { font-family: monospace; white-space: pre; }
+        .attention { background-color: #90EE90; }  /* Light green */
+        .label { background-color: #FFB6C6; }      /* Light red */
+        .token { color: #0066cc; }                 /* Blue */
+        .stats { font-weight: bold; }
+    </style>
+    """
+
+
+def escape_special_chars(text):
+    """Convert special characters to their string representation"""
+    return (
+        text.replace("\n", "\\n")
+        .replace("\t", "\\t")
+        .replace("\r", "\\r")
+        .replace(" ", "␣")
+    )  # Optional: show spaces with visible character
+
+
+def visualize_tokens_html(input_ids, attention_mask, labels, tokenizer):
+    """
+    Visualize model inputs using HTML colored text representation for Jupyter Notebook
+    with special characters shown as strings
+    """
+    # Convert to numpy if tensors
+    attention_mask = np.array(attention_mask).flatten()
+    labels = np.array(labels).flatten()
+    input_ids = np.array(input_ids).flatten()
+
+    # Decode tokens and escape special characters
+    tokens = [escape_special_chars(tokenizer.decode(id_)) for id_ in input_ids]
+
+    # Create HTML output
+    html_output = [create_color_style()]
+
+    # Header
+    html_output.append("<h3>**Token Visualization**</h3>")
+
+    # Legend
+    html_output.append(
+        """
+    <div style='margin: 10px 0;'>
+        <strong>Legend:</strong><br>
+        <span class='attention'>&nbsp;&nbsp;&nbsp;&nbsp;</span> Active Attention<br>
+        <span class='label'>&nbsp;&nbsp;&nbsp;&nbsp;</span> Label Present<br>
+        <span class='token'>Text</span> Token Text<br>
+        Special Characters: \\n (newline), \\t (tab), ␣ (space)
+    </div>
+    """
+    )
+
+    # Token alignment
+    html_output.append("<strong>Token Alignment:</strong>")
+    html_output.append("<div class='token-container'>")
+
+    # Calculate maximum token length for better alignment
+    max_token_len = max(len(str(token)) for token in tokens)
+
+    for i, (input_id, token, mask, label) in enumerate(
+        zip(input_ids, tokens, attention_mask, labels)
+    ):
+        # Pad token for alignment
+        token_text = f"{token:{max_token_len}s}"
+
+        # Create classes for styling
+        classes = []
+        if mask == 1:
+            classes.append("attention")
+        if label != -100 and label != 0:
+            classes.append("label")
+
+        class_str = f"class='{' '.join(classes)}'" if classes else ""
+
+        # Create the line
+        line = f"Position {i:3d}: <span {class_str}><span class='token'>{token_text}</span></span> "
+        line += (
+            f"(Mask: {int(mask)}, Label: {int(label)}, Inpu_id: {int(input_id)})<br>"
+        )
+        html_output.append(line)
+
+    html_output.append("</div>")
+
+    # Statistics
+    html_output.append(
+        """
+    <div class='stats' style='margin-top: 10px;'>
+        Statistics:<br>
+        Total tokens: {}<br>
+        Active attention tokens: {}<br>
+        Labeled tokens: {}<br>
+    </div>
+    """.format(
+            len(tokens), attention_mask.sum(), sum(labels != -100)
+        )
+    )
+
+    # Display the HTML
+    display(HTML("".join(html_output)))
+
+
+# Example usage:
+"""
+from transformers import AutoTokenizer
+import torch
+
+# Initialize tokenizer
+tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
+
+# Sample input with special characters
+text = "Hello,\nhow are\tyou?"
+inputs = tokenizer(text, return_tensors='pt')
+labels = torch.zeros_like(inputs['input_ids'])  # dummy labels
+
+visualize_tokens_html(
+    inputs['attention_mask'][0],
+    labels[0],
+    inputs['input_ids'][0],
+    tokenizer
+)
+"""

fusion_bench/utils/type.py

@@ -6,18 +6,20 @@ from typing_extensions import TypeAlias
 
 try:
     import torch
-    from torch import Tensor
+    from torch import Tensor, nn
 
     StateDictType: TypeAlias = Dict[str, Tensor]
+    TorchModelType = TypeVar("TorchModelType", bound=nn.Module)
+
 except ImportError:
     pass
 
 
-ModuleType = type(sys)
+PyModuleType = type(sys)
 T = TypeVar("T")
 T1 = TypeVar("T1")
 T2 = TypeVar("T2")
 T3 = TypeVar("T3")
 T4 = TypeVar("T4")
 
-__all__ = ["StateDictType", "ModuleType", "T", "T1", "T2", "T3", "T4"]
+__all__ = ["StateDictType", "PyModuleType", "TorchModelType", "T", "T1", "T2", "T3", "T4"]

{fusion_bench-0.2.6.dist-info → fusion_bench-0.2.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fusion_bench
-Version: 0.2.6
+Version: 0.2.7
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 License: MIT License

{fusion_bench-0.2.6.dist-info → fusion_bench-0.2.7.dist-info}/RECORD

@@ -1,12 +1,12 @@
 fusion_bench/__init__.py,sha256=68dF-zPvb8E2MgYnmgIJsxIHJBy1MApKeOrRZvQEVlg,421
 fusion_bench/compat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-fusion_bench/compat/method/__init__.py,sha256=yY-ILbwNVTCbor4Z7SOp0wRDbB8FqlaXo4sgF12EhQM,4823
-fusion_bench/compat/method/base_algorithm.py,sha256=Vsc9k04o6FAhu509xGYc1vZWkmegQOjqoqT7IJ8p7CA,1741
+fusion_bench/compat/method/__init__.py,sha256=KUKHpX7AfvB7fmOAlruWp0r1z17xpkI9l29PMvLWR9A,4956
+fusion_bench/compat/method/base_algorithm.py,sha256=0BG_QUtFCbfK8OmiYfEw8xaOj_G0unpqXDEJKXTCNYk,1952
 fusion_bench/compat/modelpool/AutoModelForSeq2SeqLM.py,sha256=m68BRGy4P-P9lLB10oXOBI-p58a-0FOPcrJ4r4MU32k,1100
-fusion_bench/compat/modelpool/__init__.py,sha256=C0CFrqaIKRiAvhT0PT3vM98fZwmpxL34wfb4FbeKcdo,4665
+fusion_bench/compat/modelpool/__init__.py,sha256=KD8Ddr9D7rJ5YdHEQsTuNmQ0bgQfqF4l3WNMtHmRHD8,4687
 fusion_bench/compat/modelpool/base_pool.py,sha256=1gxQENvdcOSdHmUbw-x7-X-aXtoSa1Gsys_on1ys8FM,10639
 fusion_bench/compat/modelpool/huggingface_clip_vision.py,sha256=LyIPgepNOK0qrk_EnBdlTC0ZnEkEZvPUy45cO60TiPU,6918
-fusion_bench/compat/taskpool/__init__.py,sha256=fTHd7_7EwSM2K06gUCQZ1jxxhl8T_kP0ouv70wBLhpI,3630
+fusion_bench/compat/taskpool/__init__.py,sha256=LHCRs7vrWMTtMfrqFRMmnNiSZnnZ7tZyVwXZxbi1jvQ,3651
 fusion_bench/compat/taskpool/base_pool.py,sha256=1AIZBxqUJgshq0Xo3Yo9es4b-8X8ksN1mFHxSOqnDsA,3307
 fusion_bench/compat/taskpool/clip_image_classification.py,sha256=ZYZsbsE-fPzm6yafA0p-6wcDwVGryLmtXXtuEXeQbTY,7425
 fusion_bench/compat/taskpool/flan_t5_glue_text_generation.py,sha256=O_WWr6Ivpqm-XbkaDsseDPKYcqp2jezxy-8fOrICFzY,5566
@@ -20,21 +20,26 @@ fusion_bench/dataset/image_dataset.py,sha256=MSZE_UESyRRQDwnkm2KpyIARUg9SWcwqnH4
 fusion_bench/dataset/imdb.py,sha256=YRzeq5z-Fl0aYcC2QtwEBWFkvucvpNo975jwjL5SZvs,260
 fusion_bench/dataset/nyuv2.py,sha256=2OdIEaY1ywFYMLUxCTpFcIctcBMFTq4nnoOkudSo-jI,3750
 fusion_bench/dataset/arc_agi/__init__.py,sha256=xj8BMG296qPMiL4NYs-ZwqcLJ6yT2wwbubyCbWPe91w,149
-fusion_bench/dataset/arc_agi/arc.py,sha256=AfRivFvuyumYKjlJq3LSbAzFAdHB0lY4NS8KlxhWqjU,9396
+fusion_bench/dataset/arc_agi/arc.py,sha256=EH51Sk9tR7fEMdpeaTUT49_9LtjIO8VMTzSfzk8SQEI,9568
 fusion_bench/dataset/arc_agi/arc_agi.py,sha256=SFOjp0yZrsoln4cQgWU2b-WfI39od6IE1Wof8Ee0888,11768
 fusion_bench/dataset/arc_agi/augmenters.py,sha256=yhTqyRk0_zamXRQ5Ev10xYc8Dc9D71BTSOkt856x33I,30890
 fusion_bench/dataset/arc_agi/messagers.py,sha256=E6BqF1iL68ge1m9wOJMSb2Pz6_5i9CR0HxBb7i73plE,53076
 fusion_bench/dataset/arc_agi/np_cache.py,sha256=Ec1DQFtlBdMy-f4dvGEhSr4jyVnBLQELwvX1ztxJKBs,5439
-fusion_bench/dataset/arc_agi/preprocess.py,sha256=SLmkhq76RJ8zTto5JHNFORYEr2GkbrhP81pKz1A8_BE,8523
+fusion_bench/dataset/arc_agi/preprocess.py,sha256=lQrXqV4SkhrxREgbqFAop-IwC5qaoxkKosoMO-ZHITY,8509
 fusion_bench/dataset/arc_agi/representers.py,sha256=-2eTYl-UcFW4zULDjkUrOQYv9P31nttMjc9eTJsaN0g,35852
-fusion_bench/dataset/llama/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-fusion_bench/dataset/llama/alpaca.py,sha256=sITFsghX2w0KzLwQ71KRz6rfsI2WLjuuKwt8OetvmCQ,4778
-fusion_bench/dataset/llama/collate.py,sha256=wcnt9Y2G4Isbdof3HAfe-xTbUThGo7IM0AZsn0FTmBs,1932
+fusion_bench/dataset/llama/__init__.py,sha256=p8M7G69L6bga4qLl5lvAO6SKNeUBn99kkJrAQEeOvHw,22
+fusion_bench/dataset/llama/alpaca.py,sha256=0nCQRBZzIPaMzA5VSJAsWw-nE0aVhiAQD5MGJRSrvEQ,7787
+fusion_bench/dataset/llama/collate.py,sha256=fSH-vKKCGCpPT47gchETXLF2yTCMPUE3NTE-inCdczg,3869
+fusion_bench/dataset/llama/metamathqa.py,sha256=z9InmEfWy_wWMbBORumFA2NakEznZWsDWgZzOsXYWhA,1617
 fusion_bench/dataset/llama/openai.py,sha256=_QXz6ciUTN8u4ILDowZPT3SQTes7ngkFZe1MRLFtVQ8,5012
+fusion_bench/dataset/llama/preference_700k.py,sha256=CqD0ZnM2F2Z3u70tD3VaQ2yPAHkOv75m-eiXiQKIRp0,2582
 fusion_bench/dataset/llama/sharegpt.py,sha256=8hdh_5BcxIyK0ByZoVLdhd_I06kpHffxQdaC6ezzHkM,5249
 fusion_bench/dataset/llama/squad.py,sha256=H0L0BHFzVTtkw7jfgTA8gzvZDhzsqfIALq1ip_BVwaM,4810
+fusion_bench/dataset/llama/stanford_shp.py,sha256=6ueXKnFXIBBobacU1h5WxGLZrSOtBkl6RJ6ec3Tf6UU,3763
+fusion_bench/dataset/llama/ultrachat.py,sha256=Go7WvrDAYnm184fdazHGRYLbSY6Xd7jrESyQeUJtOww,1736
 fusion_bench/dataset/llama/wikitext.py,sha256=9ZHR-nMfXRumd3o-PIj3n7B83YlVeqpGkZ2zJs2B-9Y,2883
-fusion_bench/method/__init__.py,sha256=NSBIKPSjcZbZDVuwr8srDDfntfz3jQilozRCqHPYj_w,5751
+fusion_bench/dataset/llama/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+fusion_bench/method/__init__.py,sha256=bwYq0wOxiTPCuR1GvLtqQ9Sx91mikfTTJKeSHFsSRn4,5781
 fusion_bench/method/base_algorithm.py,sha256=5dutGZfPqNhO8F8FOlo3UFR91TZu2Xj7O0pTB40JvWo,1135
 fusion_bench/method/dummy.py,sha256=hb1y6LR_geRZ5eRgGwt5zJUcHYorCeIbs5i76CvurUc,1031
 fusion_bench/method/ensemble.py,sha256=rGxvJTeorfcBuE_e0XO-0-MAc9un7ZCC46ikKGuAcN4,3077
@@ -48,9 +53,9 @@ fusion_bench/method/adamerging/clip_task_wise_adamerging.py,sha256=Tys9pDJzz5YNU
 fusion_bench/method/adamerging/entropy_loss.py,sha256=ZeVe0Hq1PaMfppLqDbB0MOscZUZRNh4CALrvt8pmQC0,736
 fusion_bench/method/adamerging/flan_t5_layer_wise_adamerging.py,sha256=osc6ueCgiS4u8KUV_sZkHGFBYC8dThnTSp4NB0wkQIg,12915
 fusion_bench/method/adamerging/gpt2_layer_wise_adamerging.py,sha256=jTGUbhJCV1pcJ5k5jVeAhmtHdbHK5LlEfBhF-86xWjY,13773
-fusion_bench/method/adamerging/layer_wise_adamerging.py,sha256=Zi5JW1V10DQNZ_0G9Ny_YH4ATP0mWuO6AjrZRwO-CxU,9432
+fusion_bench/method/adamerging/layer_wise_adamerging.py,sha256=6d1vWuyiAQDh_kLLrZixPyTAxovOjfq-2T2hgLGXCWg,9734
 fusion_bench/method/adamerging/llama_adamerging.py,sha256=DHm83VaaxxHFaeFY2qbxgO1Ub3Fiqawy4p5AqCkmEp4,13112
-fusion_bench/method/adamerging/min_norm_solvers.py,sha256=uolDwgTt8yLFuMcsQbAKO0jT6pbsG1YbSR77S40bKNs,8287
+fusion_bench/method/adamerging/min_norm_solvers.py,sha256=a7n2X0BE_YajlaUygyHV0yqW6-x5dTyZ5V0mt_Q69qE,8291
 fusion_bench/method/adamerging/task_wise_adamerging.py,sha256=tUy_P4lCn6u5srFCIyMdHs-Hc1MSge4meenK8UA25tw,6006
 fusion_bench/method/adamerging/utils.py,sha256=Yq8ovlpLJY-5MkSmpoB-_EMYG8cr6eyO-WUZTxKxMTI,432
 fusion_bench/method/analysis/__init__.py,sha256=EQzOCShS0hF958drq1yg2oSVsS0hvBznPxtTAWB9SGY,122
@@ -77,21 +82,22 @@ fusion_bench/method/fisher_merging/clip_fisher_merging.py,sha256=QCutGqjkfW3OWET
 fusion_bench/method/fisher_merging/fisher_merging.py,sha256=CPU-tJiDv9FCIBYl7Pn0zA5cdRB1Md5kWchRDlJgly0,20456
 fusion_bench/method/fisher_merging/gpt2_fisher_merging.py,sha256=LZmz41jZ5dSsAHxfOUpr3u2rlCgUPTDR7xMsIlQM-jc,7576
 fusion_bench/method/linear/__init__.py,sha256=ChfkoOEAb-rUKwpowFPel-a1hRfS8gCrbnWD-jlRbe4,283
-fusion_bench/method/linear/expo.py,sha256=jTZyI0dtYa4GbchJ78mJuaYrDHD8u3ywxyQvP46hZV8,2681
+fusion_bench/method/linear/expo.py,sha256=LCHTWlsPm1Mjhrq0mfpWLVC7skkI9ZksGduy3TxULoU,3939
 fusion_bench/method/linear/linear_interpolation.py,sha256=IONw9BPiRJouY8bE9Abfyz7qVI_1B1n8KGZa0f7Pza8,2157
 fusion_bench/method/linear/llama_expo.py,sha256=ccECjhAqcFmzOIDyZ7e_aPzTM2Kj8u2D8TJytyz18YM,8476
 fusion_bench/method/linear/simple_average_for_llama.py,sha256=7JlVrmTMmrePvNGnZNoxSuCSq2Vu7cPQzjGC3WWUXBE,2079
 fusion_bench/method/linear/task_arithmetic_for_llama.py,sha256=4SZpiTD7OzhWUXtcdK3PYdXbBGyDqiZd7oZOQ0lraN0,1963
-fusion_bench/method/lm_finetune/__init__.py,sha256=rIkKoxrqKEYkA7XIR6jyhwvUK_ebX2k6Fm1d7K1kU5g,92
+fusion_bench/method/lm_finetune/__init__.py,sha256=IFGAqXujX3Fabzl_tC6zZyOyPFJfVziL0qFtj5MVxj0,149
+fusion_bench/method/lm_finetune/bradley_terry_rm.py,sha256=ys_td1IeL3bzPTE0Cixlj2JooCaB7qseRwSDwroAk5A,18777
 fusion_bench/method/lm_finetune/causal_lm_pretrain.py,sha256=4CL9KGFsUzrt-edMfTooo4G4apzTH_57rso3DGGvKL0,219
-fusion_bench/method/lm_finetune/fullfinetune_sft.py,sha256=WoVOzFhg1PRUm8iPMYJ1g98-km3wux6nrUqnWXm27Pg,18364
-fusion_bench/method/lm_finetune/peftfinetune_sft.py,sha256=EwJJITxYBFtjsjunOlpSdo70dWeoHUYI-qIyelgW4n4,19834
+fusion_bench/method/lm_finetune/fullfinetune_sft.py,sha256=eZabmkL7QUHGxl0DX9aaCYerMVUTzH5fXQfJXiSQNEc,16226
+fusion_bench/method/lm_finetune/peftfinetune_sft.py,sha256=klZ_IDr5-1xoYvyVZwug9eyKXyxA3WZuSaML2jCH_Gw,16370
 fusion_bench/method/mixture_of_experts/__init__.py,sha256=r95iu1-3tgIUP7sWuAbLuqV7xexNYMYPZkM4_8egfp8,198
 fusion_bench/method/mixture_of_experts/mixtral_merging.py,sha256=-n1CLP1o08VyMSfaTq42kRutbw-cFDSCWHTu0iNh6ok,4237
 fusion_bench/method/mixture_of_experts/mixtral_upcycling.py,sha256=tQYAeS8MLFEfH3zDFfNZrML7lRnpGLN-HquQvjPtHNw,11208
 fusion_bench/method/pruning/__init__.py,sha256=3gtmay2bkdIAEGjpAhbY2ztMZOZLKhiJcKV3mCe2H5w,252
-fusion_bench/method/pruning/llama_magnitude_prune.py,sha256=ihHa8SNe0WGPuZqRKI_6S6gmH4ooTmeTRARGkJHcsos,6300
-fusion_bench/method/pruning/llama_random_prune.py,sha256=c-qV1iFSKZK1dES6gYsgWna1BUn58dtO0NjV1eIfJrg,4566
+fusion_bench/method/pruning/llama_magnitude_prune.py,sha256=40Gmy665S9XqIw027En6E5IlomOIcKECIRje7NDkH00,6300
+fusion_bench/method/pruning/llama_random_prune.py,sha256=EW7zfE-1a5VlPPrQ5xO1k1aqFcpPUfs5eSO_a4M1K90,4566
 fusion_bench/method/pruning/llama_wanda_prune.py,sha256=8pcg3X1yn8vfhV0lEg1fHP3oTzAc_-ixLmsZRdH5uPo,12070
 fusion_bench/method/pruning/magnitude_diff_pruning.py,sha256=nXRHW87_Nwiash-udnwR9iOaJMBDo7fPTmAwmSqsAaI,6451
 fusion_bench/method/pruning/prune_utils.py,sha256=ITWO8WtrhcOYXTcjc_fAAw7cyjvqFa6axawPr3uTT68,5882
@@ -128,6 +134,8 @@ fusion_bench/method/sparse_we_moe/sparse_clip_we_moe.py,sha256=J8iVYks-SQ93dqh6F
 fusion_bench/method/sparse_we_moe/sparse_we_moe.py,sha256=6OYgj_D_4xTtqy_guA7whQu76LQ7gv-U2cIZkXe7bIg,10479
 fusion_bench/method/sparselo/__init__.py,sha256=0Uk4Hq5b9iwc5yl2QTDwvBHUItN4V6lwhxDYQrFb724,107
 fusion_bench/method/sparselo/sparselo.py,sha256=qkfFwovdOA7-NUXtLYiV1iM9bglQJydfuL805azQ6Xc,38806
+fusion_bench/method/surgery/__init__.py,sha256=xVvJvVv9wJnVgJjZDgF2YliuSsy221AQ0KSwB7J7bjo,97
+fusion_bench/method/surgery/clip_layer_wise_adamerging_surgery.py,sha256=7qoWWYcTWpfXik21mEPBlMW5BM0Kc4nxSW1GL2hzP1o,5599
 fusion_bench/method/task_arithmetic/__init__.py,sha256=pSx_NV5Ra_6UXpyYWCi6ANQoAnEtymZt_X1dDN9wT4Y,96
 fusion_bench/method/task_arithmetic/task_arithmetic.py,sha256=1D0uuNtqyA1VS35jh6AnEVsX72HnT02THyerck_lmso,5441
 fusion_bench/method/ties_merging/__init__.py,sha256=9u9teBbdILbupr9jbwk-qCXSzssCssC5FUV2BfpyZM4,67
@@ -153,34 +161,41 @@ fusion_bench/metrics/text_to_image_generation/__init__.py,sha256=OEIxpKmyy6-3iWy
 fusion_bench/metrics/text_to_image_generation/aesthetic_scorer.py,sha256=-ZaD84ENPITh_K0Fe9OKYYoiGnPhlSE9gTbBqrtnqqA,4487
 fusion_bench/metrics/text_to_image_generation/compressibility.py,sha256=x4dNTFnAN4naChBDZBO-jUghnHAyobRVOupctKYRg1w,1656
 fusion_bench/metrics/text_to_image_generation/pickscore_scorer.py,sha256=aSWzl8k7z80Cirg5qdfkPsp3sMFEv_PjA1NJv3PPWXY,3115
-fusion_bench/mixins/__init__.py,sha256=hMxt39JDb_uIvNDtp6ZJEDmaQFwx8GId2VK2Wajw9Rg,791
-fusion_bench/mixins/clip_classification.py,sha256=devw9zTpyJsCfGCR_iKuuT9iPp1XWUqqRHRdliK6riM,8030
-fusion_bench/mixins/lightning_fabric.py,sha256=S81Bf9IDktaz2RM5T69TgiwPewUJfliLy6kd-dq3kdc,6163
+fusion_bench/mixins/__init__.py,sha256=AsUNvrHdNd6xht7-dfuVipmJuRfMNFlkgG-fn3ojt1U,892
+fusion_bench/mixins/clip_classification.py,sha256=lsrh-qZln1Am0ry_rJL47EFCvVZFRHD2YYk9u3eecs8,9933
+fusion_bench/mixins/fabric_training.py,sha256=ZmycEhCaNCgVi5oM9m0q6msxgk3quowmFvDAcvskFrg,13017
+fusion_bench/mixins/lightning_fabric.py,sha256=XE3OwV68YmJX7aR64uk9h1j9Qs6UPU5F7ciBIgBYyv4,6461
 fusion_bench/mixins/rich_live.py,sha256=j7wNgrgwfdpl6nCXZGF_2DLtNq2aqCb_52Qhe9QSltc,495
 fusion_bench/mixins/serialization.py,sha256=9W50JUcM6wgFlaE9H29mATLLVobYniSDxg94FfY25w0,4049
 fusion_bench/mixins/simple_profiler.py,sha256=UDPB8QAA3rtsSdnVgL9KMthDLBY1Rh4h8mtiquiCPp4,2106
 fusion_bench/mixins/optim/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 fusion_bench/mixins/optim/adamw_with_warmup.py,sha256=qTnRl8GVVIfaplOFBHnJFuZUbxPZRWRGHGNzm_EDhDE,1421
 fusion_bench/modelpool/PeftModelForSeq2SeqLM.py,sha256=rxPKTTWno3KAcTTEfydPpXx1b0EJa8PLbqrberweFF8,2108
-fusion_bench/modelpool/__init__.py,sha256=Ydu0bLWBA15QfHv3C5Tme0ltwHYTQFm2R-FpBVzgb-o,1266
+fusion_bench/modelpool/__init__.py,sha256=LtcCWTcYkVaz7ZxEseWRwKDk3cSTg95-WQiriXBHLSA,1401
 fusion_bench/modelpool/base_pool.py,sha256=WCpDt0MpwIk_djpgpM_CANomAEJ7Uoj78GSHEYzD_oU,9142
 fusion_bench/modelpool/huggingface_automodel.py,sha256=OJ6EyYyjNv1_Bhjn-zli-e__BJ0xVa4Fx9lhXVb-DJo,552
 fusion_bench/modelpool/huggingface_gpt2_classification.py,sha256=j8nicVwtoLXY4RPE2dcepeEB3agBKkkH-xA3yMj1czw,2014
 fusion_bench/modelpool/nyuv2_modelpool.py,sha256=btuXmYxwfjI6MnGakhoOf53Iyb9fxYH20CavGTrTcnA,1375
-fusion_bench/modelpool/causal_lm/__init__.py,sha256=kxfDP9q5yTY-SijUu7YafgV56x7c-7LBvApbKJzP3eQ,78
-fusion_bench/modelpool/causal_lm/causal_lm.py,sha256=AhuwsJLNiXUHl0Besyq2pzYo6G1_9r-iSuIIZm_70FM,4987
+fusion_bench/modelpool/causal_lm/__init__.py,sha256=F432-aDIgAbUITj4GNZS9dgUKKhaDMCbTeHB-9MecaQ,99
+fusion_bench/modelpool/causal_lm/causal_lm.py,sha256=k0eOOcFbswVgBYhM9CEXvdCRU9zVC8Gw78QaiMWzeWo,4487
 fusion_bench/modelpool/clip_vision/__init__.py,sha256=3b9gN2bWUsoA1EmpitnIMnIlX7nklxbkn4WJ0QJtS2c,43
 fusion_bench/modelpool/clip_vision/modelpool.py,sha256=qG-b3ms-q3gqcRf7J6wrTDdmtu2yb1E_A25tNOjSli8,2065
 fusion_bench/modelpool/seq2seq_lm/__init__.py,sha256=FnfSMHcwNHDQEMdB2HdK4WphQ6MufsRLUkczuALjM4Q,57
 fusion_bench/modelpool/seq2seq_lm/modelpool.py,sha256=IjLHi8qycWOA4Ul9jnqR48evgVXF_pfTKLPeL9XKP-s,2052
+fusion_bench/modelpool/seq_classification_lm/__init__.py,sha256=k-t4RetcDlbkRkPHNuyeV3pQEcJnFRjd9Wp5tFBb-G8,128
+fusion_bench/modelpool/seq_classification_lm/reward_model.py,sha256=NKf-eoei1GdU4ojKSpN5_kQwax4uUEStnlKyh8qOrNg,540
+fusion_bench/modelpool/seq_classification_lm/seq_classification_lm.py,sha256=sIKAmGJwfrNSuWtxzJ_-ME0gQksEYY2y-jVt7P82Qs0,3434
 fusion_bench/models/__init__.py,sha256=TNOEH_2yAQP51m9mdWepNEf9VGUZgDthtgXbs4rhb4M,100
-fusion_bench/models/hf_clip.py,sha256=yOQ6UKMymQ3GcfpPm26QiToPztij-cXukNMMKXTmUrw,5745
+fusion_bench/models/hf_clip.py,sha256=056UHeSjKKDYXg-o7CC2zsx4fC9R6IBkPGI8IFhWTNw,7291
 fusion_bench/models/parameter_dict.py,sha256=hRie26WIeXU-wvY6JeGaP8LvpMqbuZA6Ia_1vOFMuu4,2294
 fusion_bench/models/rankone_moe.py,sha256=uwpAqk1cwxxprQ0hxuAwRuPvHDxxBKBDahd9vcaafXs,14248
 fusion_bench/models/separate_io.py,sha256=5AJlCxkHdVVffITnIRlF3ZIaKLRWDhJESVQN1lX-ZhU,3835
 fusion_bench/models/sparse_we_moe.py,sha256=b-yIeCsl2rz0i7BP9g_fqCEam7KUNjNX_J8oyZV6MJ8,16509
-fusion_bench/models/utils.py,sha256=7HKXRiWHeoNWp8LyDemG2irnMPkT9qg2ExvxjE5mUck,1858
+fusion_bench/models/utils.py,sha256=AQFI2UZSItKfJpG8ex74FPjn_SjsADLhvpv1GYqu43U,2065
 fusion_bench/models/we_moe.py,sha256=0U-m3mhzb4vFLIzn2jd7j_SQOF9lot4ddzq0l_VPp9g,8424
+fusion_bench/models/chat_templates/__init__.py,sha256=v9vKrCfBgZ3UsMBQatZv1Z-ayPualBl5ciV0aO3p3iY,85
+fusion_bench/models/chat_templates/llama_3_Instruct.py,sha256=E6grNPECr0r1KDPIGW_DmpKQw5-Dh5WbMiTaHWDXwXo,4008
+fusion_bench/models/chat_templates/load_tokenizer.py,sha256=yRs3dB2tZo0Oh-YLJcMZzWSQ5Ps8KXrggZNb5F-aBuM,1400
 fusion_bench/models/linearized/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 fusion_bench/models/linearized/linearized_model_utils.py,sha256=5yKXReQHIwDttzT_oXwY_iIpaG1zIU0Nv93BWmmOqrg,3212
 fusion_bench/models/linearized/vision_model.py,sha256=HhbhtyoLD1qVvh1Sgl_beYF2W7AvMevmUy4Jx2XlcsY,4636
@@ -212,15 +227,21 @@ fusion_bench/models/nyuv2/resnet.py,sha256=PcCfBhEsxm7W8cu3epBbIbCYFARPrPTamIa3T
 fusion_bench/models/nyuv2/resnet_dilated.py,sha256=4EXB6vrBJS307YP6k-TRY1dFJ50LURcTuzqN4tZzYRk,3125
 fusion_bench/models/smile_moe/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 fusion_bench/models/smile_moe/linear.py,sha256=voFvx4Nnfgc6YReBcY9FUGG3WrxVRjyD3odX4jIS5Eg,8724
+fusion_bench/models/surgery/surgerymodelwrapper.py,sha256=zkiV_IF4-7CfzUND--fGCzgz-Sa-dCUz5CmVoPjQt1c,5132
 fusion_bench/models/wrappers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 fusion_bench/models/wrappers/ensemble.py,sha256=wIMZMRyXw5boWAm96c4Tiyebs_HDQovKxpGQ8rLnHUQ,6308
-fusion_bench/models/wrappers/layer_wise_fusion.py,sha256=hxyizABheJds7U_I_0h5yZh2ZvhQlkfycLJ9K5DZQ3c,12210
-fusion_bench/models/wrappers/task_wise_fusion.py,sha256=RBs1NL0eFDCMTIQg9kZXrWZ32n1RrKkcAN8OuKdN6Qw,8344
-fusion_bench/optim/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-fusion_bench/optim/mezo.py,sha256=WNcJw-Az6wnctc6pqVAloDifqgef31ZI2jwlpMKmlfo,3693
+fusion_bench/models/wrappers/layer_wise_fusion.py,sha256=tISTe__HvlaHVVKkfa0nX3JRRDYLHetJ4BzIqGq-058,12316
+fusion_bench/models/wrappers/task_wise_fusion.py,sha256=gNOU1t1JUcBr3V0Apa1uvJDl5BDM2is85lkEF1SfPRo,8404
+fusion_bench/optim/__init__.py,sha256=lemrcuiA6OLjQkpYm-RP-Ox2MgjngN1ywvCo0NgShlM,61
+fusion_bench/optim/exception.py,sha256=fMgo1heiqfGhuI5RIbf30BwWSShn5RQiyeb30QtfTI0,1607
+fusion_bench/optim/mezo.py,sha256=Vm4vMGh10Fhe28_9L1MK8r_U7DrurA8Liprh2_gn4_U,3646
+fusion_bench/optim/lr_scheduler/__init__.py,sha256=W7CsdW4XKqXbNfzjvv2wmrvNWwfH_sQ-wiBViRPlP3U,29
+fusion_bench/optim/lr_scheduler/linear_warmup.py,sha256=Dvy_TCUuAQHlbDF2jo2_502Ae4JWXGrtZL3gwA_H6ZI,6566
+fusion_bench/optim/lr_scheduler/utils/__init__.py,sha256=GfZk9VYL3cFE1Qy2xQpGc1GCgnjySk5-D7EVRZ-C05Q,29
+fusion_bench/optim/lr_scheduler/utils/visualization.py,sha256=Ea1n9ElNizAe0iUnjynyfteuZunv2-UBMN_NfEU2imA,3490
 fusion_bench/programs/__init__.py,sha256=oGoRp2TMI6ELxyfkeTg2h27hZJEDz9x31AsmvwvNvJw,508
 fusion_bench/programs/base_program.py,sha256=0dX_KcMWASo53pr-ldzfUBWIjEXy6oeDWZBrfc7FIk8,195
-fusion_bench/programs/fabric_fusion_program.py,sha256=tXUdVfLCVqjUqn8rGnc7BuXCDw0dLOkqqYtq4O6EnX4,12124
+fusion_bench/programs/fabric_fusion_program.py,sha256=Bf4lnntM1J1hxKVm4Av0ohAmSqzDxOzWg75rzqps0qE,12297
 fusion_bench/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 fusion_bench/scripts/cli.py,sha256=497nhqnJAwxkqU2WCMUqcAIvqTmGRdQaByWGNTX_onY,1131
 fusion_bench/scripts/imgui.py,sha256=r9Glbfbwu3JCsX9TKQFwcHarvwA_G7ff0jWBUPW1S1U,7613
@@ -232,12 +253,13 @@ fusion_bench/taskpool/__init__.py,sha256=_qaYgzYnvrJDrZ2DjKXMvOFbelsLrujCKa_gP3U
 fusion_bench/taskpool/base_pool.py,sha256=FaP0nndeSsrwbdd9mKa_CedbX9T5AHJmxk7Lc0NEVNY,835
 fusion_bench/taskpool/dummy.py,sha256=Di9JZO3XyDYn6wAGukrJMTnkS_NaxGTeQYo_3j1JD3Y,1675
 fusion_bench/taskpool/gpt2_text_classification.py,sha256=S4YyrcJhD4JOgvHF-AVG-gENgVGl-wpQZr1SbiThM04,4886
-fusion_bench/taskpool/nyuv2_taskpool.py,sha256=lnaR1oVm0pO2CA9EVV4uk3fiWYHD-F0GzPrUUARD75I,1970
+fusion_bench/taskpool/nyuv2_taskpool.py,sha256=Y-TI-rzh9udCjX3FJ11ZbIG7CGrjDccGc-Ch1Ug6cRY,2059
 fusion_bench/taskpool/clip_vision/__init__.py,sha256=4xGO7rRbRpXF-I34A3WEMU4vydgfdtvXQ57ThaFcpmE,214
 fusion_bench/taskpool/clip_vision/clip_rankone_moe_taskpool.py,sha256=JKbRrGaRYztgZ-P0U767HISe40UpDVQ7fn6Tf2rrug0,4891
 fusion_bench/taskpool/clip_vision/clip_sparse_wemoe_taskpool.py,sha256=hVDTtg-oXqRFmAE2wZPFpk_kvtdk_wS-2-ev2ujEJBs,5390
-fusion_bench/taskpool/clip_vision/taskpool.py,sha256=NRFXsp2N8PMQzZgFHy2yfJMjoYbDaxQpPTZ4-4EHPBY,13942
+fusion_bench/taskpool/clip_vision/taskpool.py,sha256=xbJHQXUYd2ZDs-oIyE-3knCsPdiUbZCKN7O86kPwpsQ,14907
 fusion_bench/taskpool/llama/__init__.py,sha256=iB4ESMgnsl0m-z0YtRdPZiwGGv96-86R8pbSnkdet8Q,57
+fusion_bench/taskpool/llama/reward_model.py,sha256=y5a_dNDzjYSoV9RNRdmDPhALM7i20A5EiFYSjrGTpWo,5028
 fusion_bench/taskpool/llama/test_generation.py,sha256=kJ_5GruG12FsuJHDh_S7pbQgwEojTqhGpA_wVNH5KPc,6675
 fusion_bench/tasks/__init__.py,sha256=Z_ePIp4Xizkj78QopLg1dZkJAN_IF73MkbR_nkfHQ9Y,52
 fusion_bench/tasks/base_task.py,sha256=Fg_pdZhld-2KPKX0C1WrxaTz0EYWrvJerAHO-hA03GI,412
@@ -273,7 +295,7 @@ fusion_bench/utils/data.py,sha256=51nbgOnayyerLBUGHrlm9iilGjhJsBkXKKGXOKgLRW8,61
 fusion_bench/utils/devices.py,sha256=72HeUVVlVGTt97JA7KFG3D8BM8VHqR-y1nkr9Bm-PRE,7578
 fusion_bench/utils/dtype.py,sha256=kYoEGqsXitnwOU3W7ivqhQ0OjdI7MGu1VsyMJS4cSyQ,4299
 fusion_bench/utils/functools.py,sha256=7_tYJ2WD88_2DDuOOj5aZz3cYuslYH5tsVyIgCeLtmk,1318
-fusion_bench/utils/hydra_utils.py,sha256=b-5UeDnUbMc2_NAbXVabyHeCwxfhmvAIx_aYUlzVVc8,159
+fusion_bench/utils/hydra_utils.py,sha256=TklUDKDEZlg4keI-TEZiqh4gFjr9-61Rt1RMlqkoSGk,1174
 fusion_bench/utils/instantiate.py,sha256=v8L9JDfh2YoEOFpIQIHomvBoqdboZdYIaHEATnD2gdQ,16972
 fusion_bench/utils/json.py,sha256=iNeZHFvpzbb4oX-52dX15De_dMcux7vQtAUFZqW12GA,1907
 fusion_bench/utils/lazy_imports.py,sha256=v5l9cpHXPMaz1IVBmB5oOqefYr9vA3XvP340xT7Wy18,2796
@@ -285,17 +307,20 @@ fusion_bench/utils/pylogger.py,sha256=a5tHfpEFbsdzw0vhQxt4BJ6CfTXaxyuwzoDFhyNy4K
 fusion_bench/utils/rich_utils.py,sha256=V_BjY3o8bXMp-kWfxle4cK48GGHDnKbVonZX65qbXAA,5464
 fusion_bench/utils/state_dict_arithmetic.py,sha256=dVPBkO8Te9_VANPbetV59ORAQTw7D3css_-d0lYgK4k,9062
 fusion_bench/utils/timer.py,sha256=RC2hP8JqaibdL0FnRyUCBRf4m7CXyfn5tE16zBWZ7hg,1338
-fusion_bench/utils/type.py,sha256=vS3uPx1AJz9ct-hD2Rp142yug8xcQ8YRAciAwENRKME,441
+fusion_bench/utils/type.py,sha256=4CmKDVL1bKZfAfEApq79YhQavG6ZGermpQThDwiVl5w,532
+fusion_bench/utils/plot/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+fusion_bench/utils/plot/token.py,sha256=QGmL_qX8drmWnN_VNLD_0YjKc1o_qahJE-svXVor8dU,1634
+fusion_bench/utils/plot/token_notebook.py,sha256=bsntXf46Zz_RavTxNiB9c3-KvHw7LFwyLeG0MIwOwF4,3739
 fusion_bench/utils/strenum/__init__.py,sha256=id9ORi1uXrDxhbmVxitJ1KDwLS4H3AAwFpaK5h1cQzw,8531
 fusion_bench/utils/strenum/_name_mangler.py,sha256=o11M5-bURW2RBvRTYXFQIPNeqLzburdoWLIqk8X3ydw,3397
 fusion_bench/utils/strenum/_version.py,sha256=6JQRo9LcvODbCOeVFYQb9HNJ_J9XiG_Zbn8ws2A3BV8,18466
 fusion_bench_config/README.md,sha256=Lc8YSBJ5oxf9KV5kKDivJ9LRyGuraGQPmBbgbdVA-j4,703
-fusion_bench_config/clip-vit-base-patch32_robustness_corrupted.yaml,sha256=ZAmS3hcFiWOkKxrlnZXg7moXuRX6lPi5uejqhEwIFV8,1164
-fusion_bench_config/fabric_model_fusion.yaml,sha256=68adtP49Kn7Qo9YjGSIe0ACxNFUuZ4FH9inq4U1-lbo,949
+fusion_bench_config/clip-vit-base-patch32_robustness_corrupted.yaml,sha256=GtK3VuD2FOpFHH_1Hi6tlaYpdLE5Cz0nYKP92Ss9G2Y,1164
+fusion_bench_config/fabric_model_fusion.yaml,sha256=1shmbuC0B9snkFkLErBCiroF-z7UnEHscyEmKBne7Oo,949
+fusion_bench_config/llama_full_finetune.yaml,sha256=J2qCNyX1GaPHg52juSRRIkfHSb6YJdiJOM3Bn3vAo0A,771
 fusion_bench_config/llama_magnitude_pruning.yaml,sha256=xFyDJpb8gyIjosteOpEW9eayONWhl0B763r1XmO-9w8,633
 fusion_bench_config/llama_model_fusion.yaml,sha256=EJRsSbt1zttAXAAy_-5NcMkWwhYrl0osjKGXQopu4bo,588
-fusion_bench_config/llama_weighted_average.yaml,sha256=SUP6pTcqMF_5lGgvRd_iWgmmD_s4iMGDZBRPfW38HGo,960
-fusion_bench_config/nyuv2_config.yaml,sha256=1rW-5ZsJOUCCQAvShdISmmYwyvY5vl3tukPTgtcpocY,410
+fusion_bench_config/nyuv2_config.yaml,sha256=SYaafywjOIKK1f-Nl_K5EuBjXH2oN1whlqjN_dXXO-A,492
 fusion_bench_config/nyuv2_mtl_train.yaml,sha256=RfsrboIpL9Cct2RkRrKxXAqH4jLi1NECHbwH8iOGtDY,591
 fusion_bench_config/dataset/image_classification/test/cifar10.yaml,sha256=wlzzTXAKzBjQXKYWhZPGcwydxwgeAV8sM4Dp3GBHbw0,78
 fusion_bench_config/dataset/image_classification/test/cifar100.yaml,sha256=f_tsvq5bbw2Trp3f6mokXV7hUlfYr_yuebZkEqJqIVI,79
@@ -330,6 +355,8 @@ fusion_bench_config/dataset/image_classification/val/stanford-cars.yaml,sha256=p
 fusion_bench_config/dataset/image_classification/val/sun397.yaml,sha256=cayl6FNzxOj2UBjw0ikJoQNCdN3DX10xQmcx4ouFP-0,245
 fusion_bench_config/dataset/image_classification/val/svhn.yaml,sha256=uMdEYmc406i9HdkOLzfzBiJ8pfbYLIWT1pA_UId8HMg,265
 fusion_bench_config/dataset/image_classification/val/the_eight_tasks.yaml,sha256=x1-xurkOIQtWX-gpSwXDxA5fVY30KmrarS7EKaje33M,101
+fusion_bench_config/dataset/llm_sft/alpaca_cleaned.yaml,sha256=QJK8OM-C2cZNaC3_vbRucuWrTggw69YWEtPOzmxm6bo,174
+fusion_bench_config/dataset/llm_sft/ultrachat_200k.yaml,sha256=uqOGtDu4MNAOnwCA4Qp6elcnc8X_Y7sTrwYdruAvjGU,106
 fusion_bench_config/dataset/question_answering/search_qa.yaml,sha256=u_8UTKQFnjCwbeXqx2grC2bzLDpdEQy3s3Oxip_JEoc,118
 fusion_bench_config/dataset/question_answering/test/search_qa.yaml,sha256=oB2qP5ScTTbFGp75a5VLBaQtUa9VYHkGNhNAfe-AkvE,132
 fusion_bench_config/dataset/question_answering/train/MetaMathQA.yaml,sha256=3-stubxiEKjuBQHhqS_Tc_BqGK3IOwzaAtnO4sd5SX8,90
@@ -348,9 +375,13 @@ fusion_bench_config/dataset/text_generation/train/gsm8k_question_label.yaml,sha2
 fusion_bench_config/fabric/auto.yaml,sha256=GOCIA6s_co-JpwyBj3dM-dgWUMKciD8lFRmFThUeAsA,575
 fusion_bench_config/fabric/llama_ddp.yaml,sha256=iP-3n-hehRSjmJkdQWhDb4AkLcfFa6PFq0BCrL58xso,706
 fusion_bench_config/fabric/llama_fsdp.yaml,sha256=--_G5mcyG6L3aEBNvTjvMH8D-jD0SMXGap6V8E3jH84,575
+fusion_bench_config/fabric/llama_peft_fsdp.yaml,sha256=V-iBtvSg_m2o42ERYRxlDITqeEUBoRTMrLzfVOtN8VU,580
 fusion_bench_config/fabric/loggers/csv_logger.yaml,sha256=Pv8I-xbxrpTb_fwtDiUtCAEoCZ8QYCLu2GeJNzb3Z3c,373
 fusion_bench_config/fabric/loggers/tensorboard_logger.yaml,sha256=w9ZP1i8lRYQFslzEM98PmbcFhhn5dXReSJhLOdEi-do,381
+fusion_bench_config/fabric/loggers/wandb_logger.yaml,sha256=eF4slc6QPRuMCMJVeFHNJirsGiB15WQIxNgioXNwezc,142
+fusion_bench_config/fabric/strategy/deepspeed.yaml,sha256=uHujbd2sKrgWb5YhFTWOJCZefStv6O_HHo_GylzqYbU,344
 fusion_bench_config/fabric/strategy/llama_fsdp.yaml,sha256=WBx05GFUCuEtF-H7LhlTq95VZeaIg36hqntw478qJng,307
+fusion_bench_config/fabric/strategy/llama_peft_fsdp.yaml,sha256=xoxeQ0Pp7ecZPcAX57PhQJsqRUKhqEmNc3DXmYXqx4Y,348
 fusion_bench_config/hydra/default.yaml,sha256=TT0RaUwYgfB7pKpbHgEbmuCVTB2fx2eXxvM-Xz3SQMI,241
 fusion_bench_config/hydra/help/fusion_bench_help.yaml,sha256=v8s891Cr5wyxBXGDn_VBBwwRmb0JXOL874Sl-zNoCWA,1880
 fusion_bench_config/hydra/job_logging/rich_logging.yaml,sha256=_dYGeFTCqaPrRowLXBNMXwzYhw8ns1TkQFfALwK1aCw,441
@@ -392,8 +423,9 @@ fusion_bench_config/method/linear/simple_average_for_llama.yaml,sha256=QJR5qx9z4
 fusion_bench_config/method/linear/task_arithmetic_for_llama.yaml,sha256=N7cyHm6a2QwNsV9uaJp-eZmdbs9kmdRrkxtO58QQQgM,116
 fusion_bench_config/method/linear/weighted_average.yaml,sha256=SmELszTsJU63e8KwIrPmSqKmOmH-rz42zeumQZHoVDY,187
 fusion_bench_config/method/linear/weighted_average_for_llama.yaml,sha256=r8BlNqzRfn--_gDSff6KI8FO-elWFIszZDRV7G_nvHw,499
-fusion_bench_config/method/lm_finetune/fullfinetune_sft.yaml,sha256=iJgRZiT-fic7jJOMSmq-4vslQXBIoE7IdrxPC4GQ9Cs,1157
-fusion_bench_config/method/lm_finetune/peftfinetune_sft.yaml,sha256=_LIlnNoGLJfJpchB9AYvZMRby8oG_PU3p7mdA24Eq0k,1556
+fusion_bench_config/method/lm_finetune/bradley_terry_rm.yaml,sha256=em0Lnodl9bg8dos9MODMXjKtxWCXwQArjLT2z4TC3Q0,1352
+fusion_bench_config/method/lm_finetune/fullfinetune_sft.yaml,sha256=edj3juaYos2I9oQ8J6NKQNcNwqwcQGD74ZMosDsB5SY,1341
+fusion_bench_config/method/lm_finetune/peftfinetune_sft.yaml,sha256=9S-qsWUIALRwWd_gzNF1bwIuPPGP1MmqTpdQ53cwZmc,1628
 fusion_bench_config/method/pruning/llama_magnitude_pruning.yaml,sha256=Px8LU_UtDz-YHDFfqQ7scEPOproiFOaudKVshrhCTgc,483
 fusion_bench_config/method/pruning/llama_random_pruning.yaml,sha256=0RiZS8d42PXZzwncPG8zcbnyYJ9vtfr2sOSqS8oDyT4,325
 fusion_bench_config/method/pruning/llama_wanda_pruning.yaml,sha256=qKe5yIRsmK2KUyYENENWlw1qlGet9TpDhR-E_uO7vAw,501
@@ -409,6 +441,7 @@ fusion_bench_config/method/smile_upscaling/smile_upscaling.yaml,sha256=G88mabTTn
 fusion_bench_config/method/sparselo_pruning/llama_iterative_sparselo.yaml,sha256=A_QFhwAzbzXxDkOPmXRbPTj2TBib66d3_3mkrf-Xu0k,641
 fusion_bench_config/method/sparselo_pruning/llama_pcp_sparselo.yaml,sha256=w1OWb38nW08K_hvrRMsCwmRxHWLGQfSSXg5nTiYaP8E,635
 fusion_bench_config/method/sparselo_pruning/llama_sparselo.yaml,sha256=J6vYIwqzh95-B3ekDias3FnCrVr4sig4zxpWyvz8hZ0,613
+fusion_bench_config/method/surgery/adamerging_surgery.yaml,sha256=Ne9JlJFgsRYcygBNCOBSN1ygBcLkE6I-8yusfTxyg-Y,826
 fusion_bench_config/method/trust_region/clip_task_arithmetic.yaml,sha256=mK09Ohsvj0Q6suj5qJM4DyCzRy192QBt4wjHS6W29IY,197
 fusion_bench_config/method/wemoe/sparse_weight_ensembling_moe.yaml,sha256=jiAco7M1XO0aekHFZKLKlXL_jRoCA8bgGD44Z7iB208,1001
 fusion_bench_config/method/wemoe/weight_ensembling_moe.yaml,sha256=OEv5yhyUCe5lXeT2PyXC49yrHXEM7i8SZDw6IQRDtAE,620
@@ -498,7 +531,11 @@ fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-base-patch32_svhn_and
 fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-large-patch14_TA8.yaml,sha256=-Tt_YggxkuIGT4_q5FR16zPvW2wWhGJ5LL8omxvHjvw,380
 fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-large-patch14_TA8_model_only.yaml,sha256=PrDQr04UnhAciDrdtUutx-prMxF1Cs4jrEar7uJ-1Es,238
 fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-large-patch14_individual.yaml,sha256=aSu0QUpcEZXKbL9PAUKCQAVvs5CksG1s7PPCvjTsIzA,234
+fusion_bench_config/modelpool/CausalLMPool/llama_alpaca_cleaned.yaml,sha256=BXsYH04vAUPT4Cpr8lS1px-CYYKaCTMRWuHGWgC6qE0,647
+fusion_bench_config/modelpool/CausalLMPool/llama_codealpaca.yaml,sha256=m7NDjkzFbGG8E8e_r2UUxtXjNbMFvtAcZOXyBaJOyX4,645
 fusion_bench_config/modelpool/CausalLMPool/llama_for_causallm.yaml,sha256=hbjSkVle5zpcqGDSMGaJ20CLoO0ljIXG-gtdONuaFBY,803
+fusion_bench_config/modelpool/CausalLMPool/llama_metamathqa.yaml,sha256=OOnKzmsz6iiO2jI5ZyGmCem_Pcs3a25Dveb5PLfwpUM,593
+fusion_bench_config/modelpool/CausalLMPool/llama_ultrachat.yaml,sha256=K403avKEQlK4uRhZYHbKGluCG37sMUjLRytBR3LspmI,577
 fusion_bench_config/modelpool/CausalLMPool/simle_mixtral_exp_v4.yaml,sha256=RXdm5BQmYfq9XWVli0NsQ1Xh7jD61XnhRBOSlmd9FcI,825
 fusion_bench_config/modelpool/CausalLMPool/single_llama_model.yaml,sha256=aX0rWwB-p4N94bPX1QGwqKNIWnTrkNMuF7sMAQHzjQE,549
 fusion_bench_config/modelpool/Seq2SeqLMPool/_template.yaml,sha256=mRx-Xx4s6_IBoJJRogIBW4egmqW0wi1kGVWp_YwYVvQ,233
@@ -507,6 +544,8 @@ fusion_bench_config/modelpool/Seq2SeqLMPool/flan-t5-base_glue_lora16.yaml,sha256
 fusion_bench_config/modelpool/Seq2SeqLMPool/flan-t5-base_glue_lora16_tta.yaml,sha256=GfTY343bt5YtxtUkQxSacrtQav9lT9Y-t1VIL1Chs4k,1726
 fusion_bench_config/modelpool/Seq2SeqLMPool/flan-t5-base_individual.yaml,sha256=2YBIzqYGluOT2r6dOFpUYE4Cbdd2XoHAUps-kCDxVPQ,185
 fusion_bench_config/modelpool/Seq2SeqLMPool/flan-t5-large_glue_lora16.yaml,sha256=W1y3fKY9UTTRyv7nqbIO5DESlQVfNsWlhkHJMUYh7B4,1824
+fusion_bench_config/modelpool/SeqenceClassificationModelPool/llama_preference700k.yaml,sha256=JUzGOLANW92Y_rljOOZKmwBQvWrJsko_ziayurzHSTY,880
+fusion_bench_config/modelpool/SeqenceClassificationModelPool/single_reward_model.yaml,sha256=Cg9W315FzKP3DC_-bkIyYZp1nU6UoSQ6xe-MsMx-P8A,626
 fusion_bench_config/taskpool/clip-vit-base-patch32_robustness_clean.yaml,sha256=vcU1ygptQ7nlufCEdKDWGMyi-OH4zJs55_vxG-iNHBc,541
 fusion_bench_config/taskpool/clip-vit-base-patch32_robustness_corrupted.yaml,sha256=Ged9KWmmGl29hq0gXzyG1DlryuLebDQAJIb_t5PvqiE,758
 fusion_bench_config/taskpool/clip-vit-base-patch32_svhn_and_mnist.yaml,sha256=gk_RB12EeYrEUNlZJHtZ3XKIm_LDraqE1hC_lpOEvtY,518
@@ -514,6 +553,7 @@ fusion_bench_config/taskpool/dummy.yaml,sha256=Id4Y_j7oc39qWjjEFG3qLmmMI1fGXXt34
 fusion_bench_config/taskpool/flan-t5_glue_text_generation.yaml,sha256=3MxfXiiwWJHEVgJ7aViTR7kzOV_YxXLL-fNHtnBaWN4,1002
 fusion_bench_config/taskpool/gpt-2_glue.yaml,sha256=16bw4-g08pL51M4OWAP08kWZPj6JcEefz4Xc4XhCTLQ,950
 fusion_bench_config/taskpool/nyuv2_taskpool.yaml,sha256=UaxDpFqEPkEz3h2CjFleUxsmnFnaY1aLXerkud8Zm9s,133
+fusion_bench_config/taskpool/reward_model_evaluation.yaml,sha256=qaUqKcb6E_XfoS5J-FcteMJzFzthNb4aLpB9aKW4jbU,442
 fusion_bench_config/taskpool/CLIPVisionModelTaskPool/_template.yaml,sha256=X70J8HMoIcpKaYGjg7KaaXRvz1tPUbuCqKvK9-kGHrM,1310
 fusion_bench_config/taskpool/CLIPVisionModelTaskPool/clip-vit-classification_TA8.yaml,sha256=eoNUaX-cBjpJJt0BYb-ZCNiIlv1SarX9toiGAwHbES0,227
 fusion_bench_config/taskpool/CLIPVisionModelTaskPool/clip-vit-classification_TA8_B16.yaml,sha256=AmfMrb2_wXDfRtUDsSCNgbuSicNsfC2vRlwXW-uNeJA,784
@@ -522,9 +562,9 @@ fusion_bench_config/taskpool/CLIPVisionModelTaskPool/clip-vit-classification_TA8
 fusion_bench_config/taskpool/CLIPVisionModelTaskPool/clip-vit-classification_TA8_with_control_task.yaml,sha256=9hbvC3k5x6NpA9tRDYeORhrjEyd2VH5ztMdLU67Adjk,249
 fusion_bench_config/taskpool/CLIPVisionModelTaskPool/clip_rankone_wemoe_clip-vit-classification_TA8.yaml,sha256=iQMj2VpDTe_D8OfCo94w5Ud2MON-EGa0DzVr6UmphrA,436
 fusion_bench_config/taskpool/CLIPVisionModelTaskPool/clip_sparse_wemoe_clip-vit-classification_TA8.yaml,sha256=i5Bn8bLl2cgqvrgtIGmoovUfSMehk_m-6C2wwcx5JMU,435
-fusion_bench-0.2.6.dist-info/LICENSE,sha256=nhnOJlw4CPuPVE0qvkGmxfFgHmKi-6nzXvTu8t0NUdg,1066
-fusion_bench-0.2.6.dist-info/METADATA,sha256=eExQgyXjCnwYCSMfJ3h9yH4vWaviRwNogM0OMJktUDU,13528
-fusion_bench-0.2.6.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-fusion_bench-0.2.6.dist-info/entry_points.txt,sha256=iUQ8MCJvda7HP4vYh2n1Teoapb4G9PBVYZkAfcc5SHU,116
-fusion_bench-0.2.6.dist-info/top_level.txt,sha256=BuO4TL6iHL_2yPBUX9-LlIrHRczA_BNMIFwweK0PQEI,13
-fusion_bench-0.2.6.dist-info/RECORD,,
+fusion_bench-0.2.7.dist-info/LICENSE,sha256=nhnOJlw4CPuPVE0qvkGmxfFgHmKi-6nzXvTu8t0NUdg,1066
+fusion_bench-0.2.7.dist-info/METADATA,sha256=XMvOFwwYoq1_J4Fta1kJ2J0grFb4k-I3CPA_ApmjPRM,13528
+fusion_bench-0.2.7.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+fusion_bench-0.2.7.dist-info/entry_points.txt,sha256=iUQ8MCJvda7HP4vYh2n1Teoapb4G9PBVYZkAfcc5SHU,116
+fusion_bench-0.2.7.dist-info/top_level.txt,sha256=BuO4TL6iHL_2yPBUX9-LlIrHRczA_BNMIFwweK0PQEI,13
+fusion_bench-0.2.7.dist-info/RECORD,,

fusion_bench_config/clip-vit-base-patch32_robustness_corrupted.yaml

@@ -11,7 +11,7 @@ _target_: fusion_bench.programs.FabricModelFusionProgram
 _recursive_: false
 fast_dev_run: false # Run a single batch of data to test the model or method
 # Run the script without actually running the experiment, use with `print_config=true`. 
-# You can also use `--cfg` or `-c` to  show the configuration instead of runing.
+# You can also use `--cfg` or `-c` to show the configuration instead of running.
 dry_run: false
 print_config: true # Print the configuration to the console
 merged_model_save_path: null # path to save the merged model, use "{log_dir}" to refer to the logger directory, for example `merged_model_save_path=\{log_dir\}/merged_model`

fusion_bench_config/dataset/llm_sft/alpaca_cleaned.yaml

@@ -0,0 +1,6 @@
+alpaca-cleaned:
+  _target_: fusion_bench.dataset.llama.alpaca.load_tokenized_alpaca_dataset
+  tokenizer: ???
+  path: "yahma/alpaca-cleaned"
+  split: train
+  cache_path: null

fusion_bench_config/dataset/llm_sft/ultrachat_200k.yaml

@@ -0,0 +1,3 @@
+ultrachat-200k:
+  _target_: fusion_bench.dataset.ultrachat.load_tokenized_ultrachat_200k
+  tokenizer: ???

fusion_bench_config/fabric/llama_peft_fsdp.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - loggers: tensorboard_logger
+  - strategy: llama_peft_fsdp
+  - _self_
+
+_target_: lightning.Fabric
+_recursive_: true
+# Number of devices to train on (``int``), which GPUs to train on (``list`` or ``str``), or ``"auto"``.
+# The value applies per node.
+devices: auto
+# The hardware to run on. Possible choices are:
+# ``"cpu"``, ``"cuda"``, ``"mps"``, ``"gpu"``, ``"tpu"``, ``"auto"``.
+# for example: fabric.accelerator=cpu
+accelerator: auto
+# reference to the precision policy: https://lightning.ai/docs/fabric/stable/api/fabric_args.html#precision
+precision: bf16-true

fusion_bench_config/fabric/loggers/wandb_logger.yaml

@@ -0,0 +1,2 @@
+# https://lightning.ai/docs/fabric/2.4.0/guide/loggers/wandb.html#weights-and-biases
+_target_: wandb.integration.lightning.fabric.WandbLogger

fusion_bench_config/fabric/strategy/deepspeed.yaml

@@ -0,0 +1,10 @@
+# https://lightning.ai/docs/fabric/2.4.0/api/generated/lightning.fabric.strategies.DeepSpeedStrategy.html#deepspeedstrategy
+_target_: lightning.fabric.strategies.DeepSpeedStrategy
+
+accelerator: null
+zero_optimization: true
+stage: 2
+offload_optimizer: false
+offload_parameters: false
+offload_params_device: "cpu"
+offload_optimizer_device: "cpu"

fusion_bench_config/fabric/strategy/llama_peft_fsdp.yaml

@@ -0,0 +1,9 @@
+_target_: lightning.fabric.strategies.FSDPStrategy
+sharding_strategy: FULL_SHARD
+state_dict_type: full # Save a single, consolidated checkpoint file
+cpu_offload: false
+auto_wrap_policy:
+  _target_: fusion_bench.mixins.lightning_fabric.get_size_based_auto_wrap_policy
+activation_checkpointing_policy: ${.auto_wrap_policy}
+# limit_all_gathers: true
+

fusion_bench_config/fabric_model_fusion.yaml

@@ -11,7 +11,7 @@ _target_: fusion_bench.programs.FabricModelFusionProgram
 _recursive_: false
 fast_dev_run: false # Run a single batch of data to test the model or method
 # Run the script without actually running the experiment, use with `print_config=true`. 
-# You can also use `--cfg` or `-c` to  show the configuration instead of runing.
+# You can also use `--cfg` or `-c` to show the configuration instead of running.
 dry_run: false
 print_config: true # Print the configuration to the console
 merged_model_save_path: null # path to save the merged model, use "{log_dir}" to refer to the logger directory, for example `merged_model_save_path=\{log_dir\}/merged_model`