Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,14 @@ 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
+Miniconda3-latest-Linux-x86_64.sh filter=lfs diff=lfs merge=lfs -text
+wandb/offline-run-20251225_114458-oskki2mm/run-oskki2mm.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20251225_120039-17r3f9k1/run-17r3f9k1.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20251225_121621-23cpikjp/run-23cpikjp.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20251225_124757-209hc11t/run-209hc11t.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20251225_135106-poxltdws/run-poxltdws.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20251225_144139-zgvpe4a2/run-zgvpe4a2.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20251225_170452-24ta21sy/run-24ta21sy.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20251225_200829-ad33gu8w/run-ad33gu8w.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20251225_234450-ddqothxs/run-ddqothxs.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20251226_013720-ly1ggf35/run-ly1ggf35.wandb filter=lfs diff=lfs merge=lfs -text

.vscode/launch.json

@@ -0,0 +1,25 @@
+{
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "att_matching (cifar10, GPU0)",
+      "type": "debugpy",
+      "request": "launch",
+      "program": "${workspaceFolder}/src/imagenet/att_matching.py",
+      "console": "integratedTerminal",
+      "python": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python",
+      "env": {
+        "CUDA_VISIBLE_DEVICES": "0"
+      },
+      "args": [
+        "--input-size", "32",
+        "--data-set", "CIFAR10",
+        "--model-a", "/mnt/d/NCKH/weights/lmc/cifar10/vit-finetune/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20/best_19110",
+        "--model-b", "/mnt/d/NCKH/weights/lmc/cifar10/vit-finetune/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40/best_16380",
+        "--data-path", "/mnt/d/NCKH/datasets/cifar10",
+        "--save-path", "plots/cifar10"
+      ],
+      "justMyCode": false
+    }
+  ]
+}

=5.28.0

@@ -0,0 +1 @@
+Requirement already satisfied: protobuf in /home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages (4.21.12)

LMC_of_MoE.egg-info/PKG-INFO

@@ -0,0 +1,112 @@
+Metadata-Version: 2.4
+Name: LMC-of-MoE
+Version: 0.3.1
+Summary: Linear Mode Connectivity experiments with Mixture-of-Experts models
+Author-email: Your Name <you@example.com>
+License: MIT
+Project-URL: Homepage, https://github.com/yourname/LMC-of-MoE
+Project-URL: Source, https://github.com/yourname/LMC-of-MoE
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# On Linear Mode Connectivity of Mixture-of-Experts Architectures
+
+
+[![Documentation](https://img.shields.io/badge/docs-passing-brightgreen)](https://github.com/repo/docs)
+[![Paper](https://img.shields.io/badge/arXiv-XXXX.XXXXX-blue)](https://arxiv.org/abs/XXXX.XXXXX)
+
+This repository accompanies the paper:
+***“On Linear Mode Connectivity of Mixture-of-Experts Architectures”*** (Neurips 2025 Submission)
+<p align="center"><strong>ImageNet: Linear Mode Connectivity</strong></p>
+<p align="center">
+  <img src="plots/imagenet/imagenet_lmc.png" width="500px"/>
+</p>
+
+
+## Installation
+
+```bash
+git clone https://github.com/repo/lmc-moe.git
+cd moe-lmc
+pip install -e .
+pip install -r requirements.txt
+```
+
+## Repository Structure
+
+```bash
+src/
+├── agnews/               # Appendix experiment: Reinit FFN
+├── cifar10/              # Main experiment
+├── cifar100/             # Main experiment
+├── dbpedia/              # Appendix experiment: Reinit FFN
+├── enwik8/               # Appendix experiment: Reinit FFN
+├── imagenet/             # Main experiment
+├── imdbreview/           # Appendix experiment: Reinit FFN
+├── lm1b/                 # Main experiment
+├── mnist/                # Main experiment
+├── penn/                 # Appendix experiment: Reinit FFN 
+├── transfer_learning/    # Main experiment
+├── wikitext103/          # Main experiment
+├── datasets.py
+├── utils.py
+├── weight_matching.py
+└── online_stats.py
+```
+
+Each dataset directory includes a standalone `README.md` with detailed steps for data preparation, training, and evaluation.
+
+
+## Linear Mode Connectivity Results
+
+###  ImageNet, WikiText103, One Billion Word (lm1b)
+
+
+
+<p align="center"><strong>WikiText103: Linear Mode Connectivity</strong></p>
+<p align="center">
+  <img src="plots/wikitext103/wikitext_lmc.png" width="500px"/>
+</p>
+
+<p align="center"><strong>One Billion Word (LM1B): Linear Mode Connectivity</strong></p>
+<p align="center">
+  <img src="plots/lm1b/lm1b_lmc.png" width="500px"/>
+</p>
+
+
+## Getting Started
+
+Each dataset experiment can be run individually. See the corresponding `src/<dataset>/README.md` for configuration options.
+
+
+## Citation
+
+If you find this work helpful, please consider citing:
+
+```bibtex
+@article{our2025moelmc,
+  title={On Linear Mode Connectivity of Mixture-of-Experts Architectures},
+  author={Coauthors},
+  journal={arXiv:XXXX.XXXXX},
+  year={2025}
+}
+```
+
+
+## Acknowledgements
+
+We thank contributors and maintainers of open-source libraries including PyTorch, JAX, Flax, and HuggingFace Transformers. Special thanks to the authors of recent works on LMC and MoE architectures for foundational insights.
+
+
+## Contributing
+
+We welcome pull requests and suggestions. Please ensure new features or bug fixes include tests where appropriate and follow existing code style.
+
+
+## License
+
+This project is licensed under the MIT License.
+

LMC_of_MoE.egg-info/SOURCES.txt

@@ -0,0 +1,167 @@
+.gitignore
+README.md
+password.txt
+pyproject.toml
+requirements.txt
+.vscode/launch.json
+LMC_of_MoE.egg-info/PKG-INFO
+LMC_of_MoE.egg-info/SOURCES.txt
+LMC_of_MoE.egg-info/dependency_links.txt
+LMC_of_MoE.egg-info/top_level.txt
+modal/enwik8/finetune.py
+modal/enwik8/matching.py
+modal/enwik8/__pycache__/finetune.cpython-310.pyc
+modal/imagenet/generalization.py
+modal/lm1b/finetune.py
+modal/lm1b/__pycache__/finetune.cpython-310.pyc
+modal/text8/finetune.py
+modal/text8/train_model.py
+modal/text8/__pycache__/finetune.cpython-310.pyc
+modal/text8/__pycache__/train_model.cpython-310.pyc
+plots/enwik8/[finetune-learnable-indice0,1,2,3,4,5,6,7,8,9,10,11-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0,1,2,3,4,5,6,7,8,9,10,11-heads8-shared1-routed0-topk0-seed20]_row0.pdf
+plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed20].pdf
+plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed40].pdf
+plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed40].pdf
+plots/enwik8/[finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed20].pdf
+plots/enwik8/[finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed40].pdf
+plots/enwik8/learnable-indice0-heads16-shared1-routed0-topk0.pdf
+plots/enwik8/learnable-indice0-heads4-shared1-routed0-topk0.pdf
+plots/enwik8/learnable-indice0-heads8-shared1-routed0-topk0.pdf
+plots/imagenet/.pdf.png
+plots/lm1b/learnable-indice0-heads12-shared1-routed0-topk0.pdf
+plots/lm1b/learnable-indice0-heads16-shared1-routed0-topk0.pdf
+plots/lm1b/learnable-indice0-heads8-shared1-routed0-topk0.pdf
+plots/text8/learnable-indice0-heads16-shared1-routed0-topk0.pdf
+plots/text8/learnable-indice0-heads4-shared1-routed0-topk0.pdf
+plots/text8/learnable-indice0-heads8-shared1-routed0-topk0.pdf
+plots/wt103/learnable-indice0-heads2-shared1-routed0-topk0.pdf
+plots/wt103/learnable-indice0-heads3-shared1-routed0-topk0.pdf
+plots/wt103/learnable-indice0-heads4-shared1-routed0-topk0.pdf
+results/agnews/[idx-0-shared-0-gated-1-topk-0-seed-0.flax+idx-0-shared-0-gated-1-topk-0-seed-20.flax].json
+results/dbpedia/[lr2e-05-epochs2-batch16-seed0+lr2e-05-epochs2-batch16-seed20].json
+results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json
+results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json
+results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json
+results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json
+results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json
+results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json
+results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json
+results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json
+results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json
+results/imagenet/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20.json
+results/imagenet/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40.json
+results/imdbreview/[idx-0-shared-0-gated-1-topk-0-seed-0.flax+idx-0-shared-0-gated-1-topk-0-seed-20.flax].json
+results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20].json
+results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40].json
+results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40].json
+results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json
+results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json
+results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json
+results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json
+results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json
+results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json
+results/penn/[lr0.0002-topk2-shared0-routed2-seed0+lr0.0002-topk2-shared0-routed2-seed20].json
+results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json
+results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json
+results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json
+results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json
+results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json
+results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json
+results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json
+results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json
+results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json
+results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed20].json
+results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed40].json
+results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed40].json
+results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed20].json
+results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed40].json
+results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed40].json
+results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json
+results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json
+results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json
+scripts/azure.sh
+scripts/git.sh
+scripts/modal.sh
+scripts/enwik8/att_matching.sh
+scripts/enwik8/finetune.sh
+scripts/enwik8/finetune4.sh
+scripts/enwik8/finetune5.sh
+scripts/enwik8/plot.sh
+scripts/enwik8/rope_matching.sh
+scripts/enwik8/train_model.sh
+scripts/imagenet/att_matching.sh
+scripts/imagenet/finetune0.sh
+scripts/imagenet/finetune1.sh
+scripts/imagenet/finetune2.sh
+scripts/imagenet/finetune3.sh
+scripts/imagenet/finetune4.sh
+scripts/imagenet/finetune6.sh
+scripts/imagenet/finetune7.sh
+scripts/imagenet/plot.sh
+scripts/imagenet/temp.sh
+scripts/imagenet/train_model.sh
+scripts/lm1b/att_matching.sh
+scripts/lm1b/finetune_att.sh
+scripts/lm1b/finetune_att0.sh
+scripts/lm1b/finetune_att01.sh
+scripts/lm1b/finetune_att1.sh
+scripts/lm1b/finetune_att2.sh
+scripts/lm1b/finetune_att23.sh
+scripts/lm1b/finetune_att4.sh
+scripts/lm1b/finetune_att67.sh
+scripts/lm1b/plot.sh
+scripts/lm1b/train_model.sh
+scripts/text8/att_matching.sh
+scripts/text8/finetune.sh
+scripts/text8/plot.sh
+scripts/text8/train_model.sh
+scripts/wt103/att_matching.sh
+scripts/wt103/finetune_att6.sh
+scripts/wt103/finetune_att7.sh
+scripts/wt103/plot.sh
+scripts/wt103/rope_matching.sh
+scripts/wt103/train_model.sh
+src/imagenet/att_matching.py
+src/imagenet/data.sh
+src/imagenet/datasets.py
+src/imagenet/finetune.py
+src/imagenet/generalization.py
+src/imagenet/lmc_model.py
+src/imagenet/matching_utils.py
+src/imagenet/model.py
+src/imagenet/plot.py
+src/imagenet/readme.md
+src/imagenet/test.py
+src/imagenet/train_model.py
+src/imagenet/transfer.py
+src/imagenet/transform.py
+src/imagenet/utils.py
+src/imagenet/__pycache__/datasets.cpython-310.pyc
+src/imagenet/__pycache__/engine.cpython-310.pyc
+src/imagenet/__pycache__/lmc_model.cpython-310.pyc
+src/imagenet/__pycache__/matching_utils.cpython-310.pyc
+src/imagenet/__pycache__/model.cpython-310.pyc
+src/imagenet/__pycache__/utils.cpython-310.pyc
+src/lgmodeling/att_matching.py
+src/lgmodeling/data_utils.py
+src/lgmodeling/finetune.py
+src/lgmodeling/getdata.py
+src/lgmodeling/lmc_model.py
+src/lgmodeling/matching_utils.py
+src/lgmodeling/model.py
+src/lgmodeling/plot.py
+src/lgmodeling/train_model.py
+src/lgmodeling/utils.py
+src/lgmodeling/vocabulary.py
+src/lgmodeling/__pycache__/data_utils.cpython-310.pyc
+src/lgmodeling/__pycache__/lmc_model.cpython-310.pyc
+src/lgmodeling/__pycache__/matching_utils.cpython-310.pyc
+src/lgmodeling/__pycache__/model.cpython-310.pyc
+src/lgmodeling/__pycache__/utils.cpython-310.pyc
+src/lgmodeling/__pycache__/vocabulary.cpython-310.pyc
+temp/bert.py
+temp/gpt2.py
+temp/gptj.py
+temp/qwen.py
+temp/roformer.py
+temp/vit.py

LMC_of_MoE.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

LMC_of_MoE.egg-info/top_level.txt

@@ -0,0 +1,6 @@
+modal
+plots
+results
+scripts
+src
+temp

Miniconda3-latest-Linux-x86_64.sh

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

README.md

@@ -1,3 +1,98 @@
----
-license: apache-2.0
----
+# On Linear Mode Connectivity of Mixture-of-Experts Architectures
+
+
+[![Documentation](https://img.shields.io/badge/docs-passing-brightgreen)](https://github.com/repo/docs)
+[![Paper](https://img.shields.io/badge/arXiv-XXXX.XXXXX-blue)](https://arxiv.org/abs/XXXX.XXXXX)
+
+This repository accompanies the paper:
+***“On Linear Mode Connectivity of Mixture-of-Experts Architectures”*** (Neurips 2025 Submission)
+<p align="center"><strong>ImageNet: Linear Mode Connectivity</strong></p>
+<p align="center">
+  <img src="plots/imagenet/imagenet_lmc.png" width="500px"/>
+</p>
+
+
+## Installation
+
+```bash
+git clone https://github.com/repo/lmc-moe.git
+cd moe-lmc
+pip install -e .
+pip install -r requirements.txt
+```
+
+## Repository Structure
+
+```bash
+src/
+├── agnews/               # Appendix experiment: Reinit FFN
+├── cifar10/              # Main experiment
+├── cifar100/             # Main experiment
+├── dbpedia/              # Appendix experiment: Reinit FFN
+├── enwik8/               # Appendix experiment: Reinit FFN
+├── imagenet/             # Main experiment
+├── imdbreview/           # Appendix experiment: Reinit FFN
+├── lm1b/                 # Main experiment
+├── mnist/                # Main experiment
+├── penn/                 # Appendix experiment: Reinit FFN 
+├── transfer_learning/    # Main experiment
+├── wikitext103/          # Main experiment
+├── datasets.py
+├── utils.py
+├── weight_matching.py
+└── online_stats.py
+```
+
+Each dataset directory includes a standalone `README.md` with detailed steps for data preparation, training, and evaluation.
+
+
+## Linear Mode Connectivity Results
+
+###  ImageNet, WikiText103, One Billion Word (lm1b)
+
+
+
+<p align="center"><strong>WikiText103: Linear Mode Connectivity</strong></p>
+<p align="center">
+  <img src="plots/wikitext103/wikitext_lmc.png" width="500px"/>
+</p>
+
+<p align="center"><strong>One Billion Word (LM1B): Linear Mode Connectivity</strong></p>
+<p align="center">
+  <img src="plots/lm1b/lm1b_lmc.png" width="500px"/>
+</p>
+
+
+## Getting Started
+
+Each dataset experiment can be run individually. See the corresponding `src/<dataset>/README.md` for configuration options.
+
+
+## Citation
+
+If you find this work helpful, please consider citing:
+
+```bibtex
+@article{our2025moelmc,
+  title={On Linear Mode Connectivity of Mixture-of-Experts Architectures},
+  author={Coauthors},
+  journal={arXiv:XXXX.XXXXX},
+  year={2025}
+}
+```
+
+
+## Acknowledgements
+
+We thank contributors and maintainers of open-source libraries including PyTorch, JAX, Flax, and HuggingFace Transformers. Special thanks to the authors of recent works on LMC and MoE architectures for foundational insights.
+
+
+## Contributing
+
+We welcome pull requests and suggestions. Please ensure new features or bug fixes include tests where appropriate and follow existing code style.
+
+
+## License
+
+This project is licensed under the MIT License.
+

modal/cifar/finetune.py

@@ -0,0 +1,40 @@
+from modal import App, Image, Volume
+import modal
+import json
+model_name = "Learnable"
+app = App(f"Generalization Model {model_name} with ImageNet")
+# Build image with all local dependencies added directly
+image = (
+    Image.from_registry("nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04",add_python="3.10")  # <--- This is the required 'tag'
+    .pip_install(["torch", "flax", "pandas", "tqdm","optax", "dataclasses", "argparse","matplotlib", 
+                  "scikit-learn","wandb","timm","torchvision","datasets","transformers","timm"]) 
+    .run_commands("""pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html""")
+    .add_local_file("pyproject.toml", "/root/pyproject.toml")
+    .add_local_dir("src", "/root/src")
+)
+# Shared volume for saving outputs or checkpoints
+volume1 = Volume.from_name("weights", create_if_missing=True)
+volume2 = Volume.from_name("datasets", create_if_missing=True)
+volume3 = Volume.from_name("plots", create_if_missing=True)
+volume4 = Volume.from_name("results", create_if_missing=True)
+@app.function(
+    image=image,
+    gpu="A100-40GB",
+    timeout=3600 * 24,
+    volumes={
+        "/root/weights/":volume1,"/root/datasets/":volume2,"/root/plots/":volume3,"/root/results/":volume4,
+    },
+)
+def run_command():
+    import os
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/imagenet/finetune.py --input-size 32 --data-set CIFAR10\
+        --model-path /root/weights/lmc/cifar10/lr0.005-learnable-epochs50-batch128-seed0/best_15990 \
+        --lr 5e-3 --epochs 50 --batch-size 128 --seed 0 --num-attention-heads 4 --lmc-layer-indices 0 --warmup-epochs 5\
+        --wandb-project "LMC-Attention" --wandb-group "ViT-CIFAR10-FFN" --wandb-entity "fpt-team"\
+        --data-path /root/datasets/cifar10 --save-dir /root/weights/lmc/cifar10/vit-finetune 
+    """)
+
+if __name__ == "__main__":
+    with app.run():
+        run_command.remote()

modal/cifar/train_model.py

@@ -0,0 +1,41 @@
+from modal import App, Image, Volume
+import modal
+import json
+model_name = "Learnable"
+app = App(f"Generalization Model {model_name} with ImageNet")
+# Build image with all local dependencies added directly
+image = (
+    Image.from_registry("nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04",add_python="3.10")  # <--- This is the required 'tag'
+    .pip_install(["torch", "flax", "pandas", "tqdm","optax", "dataclasses", "argparse","matplotlib", 
+                  "scikit-learn","wandb","timm","torchvision","datasets","transformers","timm"]) 
+    .run_commands("""pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html""")
+    .add_local_file("pyproject.toml", "/root/pyproject.toml")
+    .add_local_dir("src", "/root/src")
+)
+# Shared volume for saving outputs or checkpoints
+volume1 = Volume.from_name("weights", create_if_missing=True)
+volume2 = Volume.from_name("datasets", create_if_missing=True)
+volume3 = Volume.from_name("plots", create_if_missing=True)
+volume4 = Volume.from_name("results", create_if_missing=True)
+@app.function(
+    image=image,
+    gpu="A100",
+    timeout=3600 * 24,
+    volumes={
+        "/root/weights/":volume1,"/root/datasets/":volume2,"/root/plots/":volume3,"/root/results/":volume4,
+    },
+)
+def run_command():
+    import os
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/imagenet/train_model.py \
+        --input-size 32 --data-set CIFAR10 --patch-size 4 --hidden-size 128 --num-hidden-layers 6 --warmup-epochs 5\
+        --num-attention-heads 4 --intermediate-size 512 --position-embeddings "rope" --num-labels 10\
+        --lr 5e-3 --epochs 50 --batch-size 128 --seed 0 --num-shared-experts 1 --num-routed-experts 0  --topk 0\
+        --wandb-project "LMC-Attention" --wandb-group "ViT-CIFAR10-FFN" --wandb-entity "fpt-team"\
+        --save-dir /root/weights/lmc/cifar10 --data-path /root/datasets/cifar10
+    """)
+
+if __name__ == "__main__":
+    with app.run():
+        run_command.remote()

modal/enwik8/finetune.py

@@ -0,0 +1,55 @@
+from modal import App, Image, Volume
+import modal
+import json
+model_name = "Attention Full Layers"
+app = App(f"Finetune Enwik8 {model_name}")
+# Build image with all local dependencies added directly
+image = (
+    Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04",add_python="3.10")  # <--- This is the required 'tag'
+    .pip_install(["torch", "flax", "pandas", "tqdm","optax", "dataclasses", "argparse","matplotlib", 
+                  "scikit-learn","wandb","timm","torchvision","datasets","transformers","timm"]) 
+    .run_commands("""pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html""")
+    .add_local_file("pyproject.toml", "/root/pyproject.toml")
+    .add_local_dir("src", "/root/src")
+)
+# Shared volume for saving outputs or checkpoints
+volume1 = Volume.from_name("weights", create_if_missing=True)
+volume2 = Volume.from_name("dataset", create_if_missing=True)
+volume3 = Volume.from_name("plots", create_if_missing=True)
+volume4 = Volume.from_name("results", create_if_missing=True)
+@app.function(
+    image=image,
+    gpu="H100:1",
+    timeout=3600 * 24,
+    volumes={"/root/weights/":volume1,"/root/dataset/":volume2,"/root/plots/":volume3,"/root/results/":volume4},
+)
+def run_command():
+    import os
+    os.system("pip install -e .")
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/enwik8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+        --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 16 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset enwik8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/enwik8/gpt2-finetune --data-path /root/dataset/enwik8
+    """)
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/enwik8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+        --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 16 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset enwik8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/enwik8/gpt2-finetune --data-path /root/dataset/enwik8
+    """)
+    # os.system("""
+    # WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+    #     --model-path /root/weights/enwik8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+    #     --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \
+    #     --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset enwik8 \
+    #     --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+    #     --model-save-dir /root/weights/enwik8/gpt2-finetune --data-path /root/dataset/enwik8
+    # """)
+if __name__ == "__main__":
+    with app.run():
+        run_command.remote()

modal/imagenet/generalization.py

@@ -0,0 +1,57 @@
+from modal import App, Image, Volume
+import modal
+import json
+model_name = "Learnable"
+app = App(f"Generalization Model {model_name} with ImageNet")
+# Build image with all local dependencies added directly
+image = (
+    Image.from_registry("nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04",add_python="3.10")  # <--- This is the required 'tag'
+    .pip_install(["torch", "flax", "pandas", "tqdm","optax", "dataclasses", "argparse","matplotlib", 
+                  "scikit-learn","wandb","timm","torchvision","datasets","transformers","timm"]) 
+    .run_commands("""pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html""")
+    .add_local_file("pyproject.toml", "/root/pyproject.toml")
+    .add_local_dir("src", "/root/src")
+)
+# Shared volume for saving outputs or checkpoints
+volume1 = Volume.from_name("weights", create_if_missing=True)
+volume2 = Volume.from_name("datasets", create_if_missing=True)
+volume3 = Volume.from_name("plots", create_if_missing=True)
+volume4 = Volume.from_name("results", create_if_missing=True)
+volume5 = Volume.from_name("shot_noise", create_if_missing=True)
+@app.function(
+    image=image,
+    gpu="H100",
+    timeout=3600 * 24,
+    volumes={
+        "/root/weights":volume1,"/root/datasets/":volume2,"/root/plots/":volume3,"/root/results/":volume4,
+        "/root/shot_noise":volume5,
+    },
+)
+def run_command():
+    import os
+    # os.system("pip install -e .")
+    os.system("ls /root/shot_noise")
+    # os.system("""
+    # CUDA_VISIBLE_DEVICES=0  python src/imagenet/generalization.py \
+    #     --model-a /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed0/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0/best_65052\
+    #     --model-b /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed20/best_5004\
+    #     --data-original /root/datasets/imagenet\
+    #     --data-generalization /root/shot_noise/shot_noise/1
+    # """)
+    os.system("""
+    CUDA_VISIBLE_DEVICES=0  python src/imagenet/generalization.py \
+        --model-a /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed0/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0/best_65052\
+        --model-b /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed40/best_5004\
+        --data-original /root/datasets/imagenet\
+        --data-generalization /root/shot_noise/shot_noise/4
+    """)
+    # os.system("""
+    # CUDA_VISIBLE_DEVICES=0  python src/imagenet/generalization.py \
+    #     --model-a /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed20/best_5004\
+    #     --model-b /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed40/best_5004\
+    #     --data-original /root/datasets/imagenet\
+    #     --data-generalization /root/shot_noise/shot_noise/1
+    # """)
+if __name__ == "__main__":
+    with app.run():
+        run_command.remote()

modal/imagenet/shot_noise.py

@@ -0,0 +1,65 @@
+from modal import App, Image, Volume
+import modal
+import json
+model_name = "Learnable"
+app = App(f"Generalization Model {model_name} with ImageNet")
+# Build image with all local dependencies added directly
+image = (
+    Image.from_registry("nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04",add_python="3.10")  # <--- This is the required 'tag'
+    .pip_install(["torch", "flax", "pandas", "tqdm","optax", "dataclasses", "argparse","matplotlib", 
+                  "scikit-learn","wandb","timm","torchvision","datasets","transformers","timm"]) 
+    .run_commands("""pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html""")
+    .add_local_file("pyproject.toml", "/root/pyproject.toml")
+    .add_local_dir("src", "/root/src")
+)
+# Shared volume for saving outputs or checkpoints
+volume1 = Volume.from_name("weights", create_if_missing=True)
+volume2 = Volume.from_name("datasets", create_if_missing=True)
+volume3 = Volume.from_name("plots", create_if_missing=True)
+volume4 = Volume.from_name("results", create_if_missing=True)
+volume5 = Volume.from_name("shot_noise", create_if_missing=True)
+@app.function(
+    image=image,
+    gpu="H100",
+    timeout=3600 * 24,
+    volumes={
+        "/root/weights/":volume1,"/root/datasets/":volume2,"/root/plots/":volume3,"/root/results/":volume4,
+        "/root/shot_noise":volume5,
+    },
+)
+def run_command():
+    import os
+    # os.system("pip install -e .")
+    os.system("ls /root/shot_noise")
+    # os.system("""
+    # CUDA_VISIBLE_DEVICES=0  python src/imagenet/generalization.py \
+    #     --model-a /root/weights/vit-finetune/finetune-learnable-indice0,1,2,3,4,5,6,7,8,9,10,11-heads12-shared1-routed0-topk0-mlpFalse-seed0/best_375300\
+    #     --model-b /root/weights/vit-finetune/finetune-learnable-indice0,1,2,3,4,5,6,7,8,9,10,11-heads12-shared1-routed0-topk0-mlpFalse-seed20/best_335268\
+    #     --data-original /root/datasets/imagenet\
+    #     --data-generalization /root/shot_noise/shot_noise/3
+    # """)
+    # os.system("""
+    # CUDA_VISIBLE_DEVICES=0  python src/imagenet/generalization.py \
+    #     --model-a /root/weights/vit-finetune/finetune-learnable-indice0,1,2,3,4,5,6,7,8,9,10,11-heads12-shared1-routed0-topk0-mlpFalse-seed0/best_375300\
+    #     --model-b /root/weights/vit-finetune/finetune-learnable-indice0,1,2,3,4,5,6,7,8,9,10,11-heads12-shared1-routed0-topk0-mlpFalse-seed40/best_460368\
+    #     --data-original /root/datasets/imagenet\
+    #     --data-generalization /root/shot_noise/shot_noise/3\
+    # """)
+    os.system("""
+    CUDA_VISIBLE_DEVICES=0  python src/imagenet/generalization.py \
+        --model-a /root/weights/vit-finetune/finetune-learnable-indice0,1,2,3,4,5,6,7,8,9,10,11-heads12-shared1-routed0-topk0-mlpFalse-seed20/best_335268\
+        --model-b /root/weights/vit-finetune/finetune-learnable-indice0,1,2,3,4,5,6,7,8,9,10,11-heads12-shared1-routed0-topk0-mlpFalse-seed40/best_460368\
+        --data-original /root/datasets/imagenet\
+        --data-generalization /root/shot_noise/shot_noise/3\
+    """)
+
+
+
+
+
+    
+
+
+if __name__ == "__main__":
+    with app.run():
+        run_command.remote()

modal/lm1b/finetune.py

@@ -0,0 +1,47 @@
+from modal import App, Image, Volume
+import modal
+import json
+model_name = "Learnable"
+app = App(f"Train Model {model_name} with lm1b")
+# Build image with all local dependencies added directly
+image = (
+    Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04",add_python="3.10")  # <--- This is the required 'tag'
+    .pip_install(["torch", "flax", "pandas", "tqdm","optax", "dataclasses", "argparse","matplotlib", 
+                  "scikit-learn","wandb","timm","torchvision","datasets","transformers","timm"]) 
+    .run_commands("""pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html""")
+    .add_local_file("pyproject.toml", "/root/pyproject.toml")
+    .add_local_dir("src", "/root/src")
+)
+# Shared volume for saving outputs or checkpoints
+volume1 = Volume.from_name("weights", create_if_missing=True)
+volume2 = Volume.from_name("datasets", create_if_missing=True)
+volume3 = Volume.from_name("plots", create_if_missing=True)
+volume4 = Volume.from_name("results", create_if_missing=True)
+@app.function(
+    image=image,
+    gpu="H100:2",
+    timeout=3600 * 24,
+    volumes={"/root/weights/":volume1,"/root/datasets/":volume2,"/root/plots/":volume3,"/root/results/":volume4},
+)
+def run_command():
+    import os
+    os.system("pip install -e .")
+
+    # ---------- Env “exports” to tame XLA autotuner miscompares ----------
+    # Safer kernels + determinism; disable TF32 for full fp32 matmuls
+    os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
+    os.environ["JAX_DEFAULT_MATMUL_PRECISION"] = "highest"  # same as jax.config.update(...)
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+    # --------------------------------------------------------------------
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \
+        --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 12 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\
+        --learning-rate 0.00025 --batch-size 96 --max_step 100000 --warmup_step 1000 --dataset lm1b \
+        --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/lm1b/gpt2-finetune --data-path /root/datasets/lm1b
+    """)
+
+if __name__ == "__main__":
+    with app.run():
+        run_command.remote()

modal/text8/finetune.py

@@ -0,0 +1,107 @@
+from modal import App, Image, Volume
+import modal
+import json
+model_name = "Learnable"
+app = App(f"Train Model {model_name} with Text8")
+# Build image with all local dependencies added directly
+image = (
+    Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04",add_python="3.10")  # <--- This is the required 'tag'
+    .pip_install(["torch", "flax", "pandas", "tqdm","optax", "dataclasses", "argparse","matplotlib", 
+                  "scikit-learn","wandb","timm","torchvision","datasets","transformers","timm"]) 
+    .run_commands("""pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html""")
+    .add_local_file("pyproject.toml", "/root/pyproject.toml")
+    .add_local_dir("src", "/root/src")
+)
+# Shared volume for saving outputs or checkpoints
+volume1 = Volume.from_name("weights", create_if_missing=True)
+volume2 = Volume.from_name("datasets", create_if_missing=True)
+volume3 = Volume.from_name("plots", create_if_missing=True)
+volume4 = Volume.from_name("results", create_if_missing=True)
+@app.function(
+    image=image,
+    gpu="H100:1",
+    timeout=3600 * 24,
+    volumes={"/root/weights/":volume1,"/root/datasets/":volume2,"/root/plots/":volume3,"/root/results/":volume4},
+)
+def run_command():
+    import os
+    os.system("pip install -e .")
+    os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
+    os.environ["JAX_DEFAULT_MATMUL_PRECISION"] = "highest" 
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+        --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 8  --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8
+    """)
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+        --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 8  --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8
+    """)
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+        --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 8  --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8
+    """)
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+        --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 16  --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8
+    """)
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+        --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 16  --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8
+    """)
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+        --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 16  --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8
+    """)
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+        --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 4  --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8
+    """)
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+        --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 4  --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8
+    """)
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/finetune.py \
+        --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \
+        --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128  --n_head 4  --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8
+    """)
+
+if __name__ == "__main__":
+    with app.run():
+        run_command.remote()

modal/text8/train_model.py

@@ -0,0 +1,41 @@
+from modal import App, Image, Volume
+import modal
+import json
+model_name = "Learnable"
+app = App(f"Train Model {model_name} with Text8")
+# Build image with all local dependencies added directly
+image = (
+    Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04",add_python="3.10")  # <--- This is the required 'tag'
+    .pip_install(["torch", "flax", "pandas", "tqdm","optax", "dataclasses", "argparse","matplotlib", 
+                  "scikit-learn","wandb","timm","torchvision","datasets","transformers","timm"]) 
+    .run_commands("""pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html""")
+    .add_local_file("pyproject.toml", "/root/pyproject.toml")
+    .add_local_dir("src", "/root/src")
+)
+# Shared volume for saving outputs or checkpoints
+volume1 = Volume.from_name("weights", create_if_missing=True)
+volume2 = Volume.from_name("datasets", create_if_missing=True)
+volume3 = Volume.from_name("plots", create_if_missing=True)
+volume4 = Volume.from_name("results", create_if_missing=True)
+@app.function(
+    image=image,
+    gpu="H100:1",
+    timeout=3600 * 24,
+    volumes={"/root/weights/":volume1,"/root/datasets/":volume2,"/root/plots/":volume3,"/root/results/":volume4},
+)
+def run_command():
+    import os
+    os.system("pip install -e .")
+    os.system("""
+    WANDB_MODE=online CUDA_VISIBLE_DEVICES=0  python src/lgmodeling/train_model.py \
+        --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "learnable" \
+        --num-shared-experts 1 --num-routed-experts 0  --topk 0 --rotary-dim 64 \
+        --n_layer 12 --n_embd 512 --n_head 8 --n_inner 2048 --attention-bias \
+        --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \
+        --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\
+        --model-save-dir /root/weights/text8 --data-path /root/datasets/text8
+    """)
+
+if __name__ == "__main__":
+    with app.run():
+        run_command.remote()

password.txt

@@ -0,0 +1,2 @@
+lmc-transformer: 6niAInCZxocweP6weMTKC8cobrLVHZmttuVh96OmamjSL9gjkHSdJQQJ99BFACAAAAAubb37AAASAZDO2upd
+UKayDJX6eG

pyproject.toml

@@ -0,0 +1,29 @@
+[build-system]
+requires = ["setuptools", "setuptools-scm"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "LMC-of-MoE"
+version = "0.3.1"
+description = "Linear Mode Connectivity experiments with Mixture-of-Experts models"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+authors = [
+    { name = "Your Name", email = "you@example.com" }
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/yourname/LMC-of-MoE"
+"Source" = "https://github.com/yourname/LMC-of-MoE"
+
+[tool.setuptools.packages.find]
+where = ["./"]
+
+[tool.setuptools.dynamic]
+dependencies = { file = ["requirements.txt"] }