diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..5afc1eb8cac544aa6a50b833fa195f9502e5e509 100644 --- a/.gitattributes +++ b/.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 diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000000000000000000000000000000000000..d0f0c58b00c5ce1a782e557138ffa0d7e58e1ee2 --- /dev/null +++ b/.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 + } + ] +} diff --git a/=5.28.0 b/=5.28.0 new file mode 100644 index 0000000000000000000000000000000000000000..45be4001c48426ede7af3c231659ccdc0b9f08a4 --- /dev/null +++ b/=5.28.0 @@ -0,0 +1 @@ +Requirement already satisfied: protobuf in /home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages (4.21.12) diff --git a/LMC_of_MoE.egg-info/PKG-INFO b/LMC_of_MoE.egg-info/PKG-INFO new file mode 100644 index 0000000000000000000000000000000000000000..d16853880ef0fd90ca0f43512a9b0a01f1664305 --- /dev/null +++ b/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 +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) +

ImageNet: Linear Mode Connectivity

+

+ +

+ + +## 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) + + + +

WikiText103: Linear Mode Connectivity

+

+ +

+ +

One Billion Word (LM1B): Linear Mode Connectivity

+

+ +

+ + +## Getting Started + +Each dataset experiment can be run individually. See the corresponding `src//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. + diff --git a/LMC_of_MoE.egg-info/SOURCES.txt b/LMC_of_MoE.egg-info/SOURCES.txt new file mode 100644 index 0000000000000000000000000000000000000000..60538ba4413ff8117c46a72d38e7ad1770344910 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/LMC_of_MoE.egg-info/dependency_links.txt b/LMC_of_MoE.egg-info/dependency_links.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/LMC_of_MoE.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/LMC_of_MoE.egg-info/top_level.txt b/LMC_of_MoE.egg-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..4f9f810b1c2a0a3db9f57a5824cdcb8ca8e39b91 --- /dev/null +++ b/LMC_of_MoE.egg-info/top_level.txt @@ -0,0 +1,6 @@ +modal +plots +results +scripts +src +temp diff --git a/Miniconda3-latest-Linux-x86_64.sh b/Miniconda3-latest-Linux-x86_64.sh new file mode 100644 index 0000000000000000000000000000000000000000..427f07231125aa7e15bee73e98ed88822aa04252 --- /dev/null +++ b/Miniconda3-latest-Linux-x86_64.sh @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e0b10e050e8928e2eb9aad2c522ee3b5d31d30048b8a9997663a8a460d538cef +size 156772981 diff --git a/README.md b/README.md index 7b95401dc46245ac339fc25059d4a56d90b4cde5..84f15e114b24f4045f89ae147af4985235ccc9f5 100644 --- a/README.md +++ b/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) +

ImageNet: Linear Mode Connectivity

+

+ +

+ + +## 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) + + + +

WikiText103: Linear Mode Connectivity

+

+ +

+ +

One Billion Word (LM1B): Linear Mode Connectivity

+

+ +

+ + +## Getting Started + +Each dataset experiment can be run individually. See the corresponding `src//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. + diff --git a/modal/cifar/__pycache__/finetune.cpython-313.pyc b/modal/cifar/__pycache__/finetune.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a861b1034dcf26301cd5f2229565264e2a310dd Binary files /dev/null and b/modal/cifar/__pycache__/finetune.cpython-313.pyc differ diff --git a/modal/cifar/__pycache__/train_model.cpython-313.pyc b/modal/cifar/__pycache__/train_model.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..303190051024bd08fe3fa5649cd4a7d00b73aa14 Binary files /dev/null and b/modal/cifar/__pycache__/train_model.cpython-313.pyc differ diff --git a/modal/cifar/finetune.py b/modal/cifar/finetune.py new file mode 100644 index 0000000000000000000000000000000000000000..a79c0980a953d598c11609c6a485d8115fab2741 --- /dev/null +++ b/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() diff --git a/modal/cifar/train_model.py b/modal/cifar/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..f34c866b4d57acc958b5a151105da57fb89c7486 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/modal/enwik8/__pycache__/finetune.cpython-310.pyc b/modal/enwik8/__pycache__/finetune.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84224a862c666145ebff05cf54dcc983df448899 Binary files /dev/null and b/modal/enwik8/__pycache__/finetune.cpython-310.pyc differ diff --git a/modal/enwik8/finetune.py b/modal/enwik8/finetune.py new file mode 100644 index 0000000000000000000000000000000000000000..d0aaef106b44745f3332c8fec7cd628c64af22e4 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/modal/enwik8/matching.py b/modal/enwik8/matching.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/modal/imagenet/__pycache__/generalization.cpython-313.pyc b/modal/imagenet/__pycache__/generalization.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b5f1b5d00e94d90350b1f0ba40a242e880e561e Binary files /dev/null and b/modal/imagenet/__pycache__/generalization.cpython-313.pyc differ diff --git a/modal/imagenet/__pycache__/shot_noise.cpython-313.pyc b/modal/imagenet/__pycache__/shot_noise.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1408b01130e7996f84db7fcb2c29f9676233127b Binary files /dev/null and b/modal/imagenet/__pycache__/shot_noise.cpython-313.pyc differ diff --git a/modal/imagenet/generalization.py b/modal/imagenet/generalization.py new file mode 100644 index 0000000000000000000000000000000000000000..34b9bfd256c4694896b4880ca39ea901dc55edbf --- /dev/null +++ b/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() \ No newline at end of file diff --git a/modal/imagenet/shot_noise.py b/modal/imagenet/shot_noise.py new file mode 100644 index 0000000000000000000000000000000000000000..61c4ffd9a1b2febfb0e44db83a855bfb5c94045d --- /dev/null +++ b/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() \ No newline at end of file diff --git a/modal/lm1b/__pycache__/finetune.cpython-310.pyc b/modal/lm1b/__pycache__/finetune.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ec91872f2e2c41b912b7156251fcb0bc12e59c83 Binary files /dev/null and b/modal/lm1b/__pycache__/finetune.cpython-310.pyc differ diff --git a/modal/lm1b/finetune.py b/modal/lm1b/finetune.py new file mode 100644 index 0000000000000000000000000000000000000000..a074baffdf0b163d70eb9830ca1e70cac0acd6e7 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/modal/text8/__pycache__/finetune.cpython-310.pyc b/modal/text8/__pycache__/finetune.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f67faf7646b27f219fbfba81d0ca3572f731dd8 Binary files /dev/null and b/modal/text8/__pycache__/finetune.cpython-310.pyc differ diff --git a/modal/text8/__pycache__/train_model.cpython-310.pyc b/modal/text8/__pycache__/train_model.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c7a8b26779ec4051f65d14dddbb732e5dbb774f Binary files /dev/null and b/modal/text8/__pycache__/train_model.cpython-310.pyc differ diff --git a/modal/text8/finetune.py b/modal/text8/finetune.py new file mode 100644 index 0000000000000000000000000000000000000000..1555e499694d4a63c7f241fb5f0b53dfa1214ff3 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/modal/text8/train_model.py b/modal/text8/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..fd886a4bf503900a5b3752e082c4d4298d1e7b57 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/output.txt b/output.txt new file mode 100644 index 0000000000000000000000000000000000000000..e4dc72927d3c138ab1ce9d0320389a426021c24f --- /dev/null +++ b/output.txt @@ -0,0 +1,6512 @@ +FrozenDict({ + classifier: { + bias: Array([ 0.13959047, -0.15330061, 0.23963794, 0.26991433, 0.14193657, + -0.07303967, -0.15755221, -0.00429095, -0.1597117 , 0.01030699], dtype=float32), + kernel: Array([[-0.21733038, -0.04500001, 0.6502397 , ..., -0.16637883, + -0.19785312, -0.19238858], + [-0.05325412, 0.08365323, 0.04513923, ..., 0.04858669, + -0.2313862 , 0.15936573], + [ 0.15941077, -0.5403379 , 0.17727152, ..., 0.23605543, + -0.5477616 , 0.400377 ], + ..., + [ 0.14339526, -0.40216786, -0.11720804, ..., -0.4565188 , + 0.55370927, -0.3628201 ], + [ 0.1330832 , -0.05788376, 0.05994943, ..., -0.05280861, + -0.35830075, -0.07267395], + [ 0.07371017, 0.0797717 , 0.05526711, ..., -0.15986513, + -0.19333237, -0.11412273]], dtype=float32), + }, + vit: { + embeddings: { + cls_token: Array([[[ 4.21749912e-02, -7.63344541e-02, -9.38936546e-02, + 3.23799951e-03, 1.69068575e-01, 2.25809757e-02, + 3.44559725e-04, -3.61712719e-03, -8.03779662e-02, + 3.50343175e-02, -2.25252569e-01, 8.41705650e-02, + -1.40322876e+00, -6.44051209e-02, 3.56393009e-02, + 1.55636281e-01, 1.44712580e-02, -1.67391926e-01, + 1.13313071e-01, 2.03049585e-01, -1.21489443e-01, + 3.01832557e-01, -6.39758036e-02, 2.43316695e-01, + 8.19623619e-02, 3.64291817e-01, 2.20934585e-01, + 4.44887951e-02, -1.20811634e-01, 7.45030791e-02, + 2.88756996e-01, 2.45113447e-01, -1.02346905e-01, + -1.44495424e-02, -2.92423535e-02, 1.27275869e-01, + -4.68343794e-02, 6.18743040e-02, 1.62268832e-01, + 7.68043697e-02, 8.22397470e-02, 4.15349565e-02, + 4.29975651e-02, 3.34680080e-01, 1.23265073e-01, + -1.82213470e-01, -2.61384398e-01, -2.86995154e-02, + 4.78795953e-02, -7.12198839e-02, -9.28327814e-02, + -1.04730994e-01, -2.08872005e-01, 5.08915447e-02, + -2.92177856e-01, -4.63173315e-02, -1.39276087e-01, + -6.45741746e-02, 4.31298822e-01, -2.79950589e-01, + -3.42224717e-01, -2.57990032e-01, 7.88487419e-02, + 1.75719306e-01, -1.91833928e-01, -1.60736710e-01, + 2.24118829e-01, -1.42223332e-02, -1.56884298e-01, + -1.34610429e-01, -7.02320859e-02, -5.77486083e-02, + -2.67209888e-01, -1.95046023e-01, 1.72840059e-01, + 5.77186942e-02, 1.49908558e-01, 1.20706670e-01, + 2.60736823e-01, -4.85861525e-02, 1.50602728e-01, + -1.42758638e-01, 2.74297837e-02, 2.43728027e-01, + 2.83344448e-01, -3.49160563e-03, -1.21281445e-01, + -1.72878504e-01, -1.20296307e-01, -5.02357781e-02, + 8.33122134e-02, -1.96546650e+00, 1.50719002e-01, + 2.28078086e-02, -1.70492530e-01, 4.20110017e-01, + -9.01997983e-02, -3.95133868e-02, 1.77724794e-01, + 1.20884374e-01, -3.86571921e-02, -7.12739974e-02, + 1.50522381e-01, 1.34570688e-01, 3.98968756e-01, + -1.69036195e-01, 7.98634291e-02, -2.08218116e-02, + -1.60120368e-01, -2.48059444e-02, -1.43607566e-02, + 3.75894368e-01, 8.90320763e-02, -1.83609471e-01, + 3.58941674e-01, -1.86423570e-01, 2.57358737e-02, + 1.14504859e-01, -1.51107222e-01, -1.81861663e+00, + -1.02884192e-02, -3.45049873e-02, 7.34117031e-02, + 6.87620267e-02, 6.22403771e-02, 3.00719798e-01, + 6.51065707e-02, 3.67412001e-01]]], dtype=float32), + patch_embeddings: { + projection: { + bias: Array([-0.23706354, -1.6540319 , -0.56307954, 0.26430997, 0.45925793, + 0.10770348, -0.64220166, -0.256316 , -0.70986354, 0.00280434, + -0.8621466 , 0.33027068, -0.09428401, -0.18285578, 0.48031166, + 0.31876177, 0.1595892 , -0.7937292 , 0.41048172, 0.7553594 , + -0.11688971, 1.0029606 , -0.19344132, 0.9868722 , 0.33903363, + 1.4018133 , 0.6315837 , 0.15693058, -0.43901056, -0.03160467, + 1.4104275 , 0.8133011 , -0.8685446 , 0.09167413, -0.38324404, + 0.4127237 , -0.0992351 , -0.8858117 , 0.49157575, 1.1184502 , + 0.29291674, 0.09339268, 0.10110624, 1.3492149 , 0.23373768, + -0.62969214, -1.1953127 , -0.32950965, 0.16054444, -0.04147059, + -0.32538444, -0.44066548, -0.81780577, 0.3041742 , -0.88412094, + -0.2175792 , -0.7703008 , -0.29075882, 1.4584533 , -0.90483326, + -1.0301116 , -1.9739118 , 0.30468434, 0.47119713, -2.507297 , + -0.4836628 , 0.7882562 , 0.13242652, -0.64575255, -0.5046555 , + -0.2928836 , -0.19443981, -1.8653952 , -0.6447174 , 0.53096473, + 0.53206563, 0.53241354, 0.5417067 , 0.87263006, -0.22854666, + 0.57786745, -0.3563041 , -0.03160271, 0.5421835 , 1.0412699 , + -0.1696869 , -0.4523725 , -1.6923577 , -0.46177727, -0.20988323, + 0.2240901 , -0.9807764 , 0.61122894, -1.4953612 , -0.6922171 , + 1.6215447 , -0.32785478, -0.6734319 , 0.5917648 , 0.33969748, + -0.27886963, -0.50457156, 0.5443144 , 0.5749297 , 2.0102262 , + 0.15956564, 0.03995752, 0.18428732, 0.1132739 , 1.1407704 , + 0.10695937, 0.6835315 , 0.15768398, -0.5695771 , 0.45638716, + -0.6475178 , -0.24609064, 0.36820015, -0.7532235 , -0.60841304, + -0.13496049, -0.49064922, 0.31223625, 0.27192888, 0.27554023, + 1.1431948 , 0.30906013, 1.071581 ], dtype=float32), + kernel: Array([[[[ 0.1773554 , -0.0076261 , 0.39377195, ..., 0.09380915, + -0.01814741, 0.2283536 ], + [-0.2979896 , 0.02483123, 0.03855339, ..., 0.09292422, + 0.0840548 , 0.02128509], + [-0.11292487, -0.09941795, -0.21962844, ..., -0.00123785, + 0.14886236, 0.08236729]], + + [[ 0.32237047, -0.10110885, 0.5543121 , ..., 0.20536377, + -0.33312383, 0.189913 ], + [-0.23798265, -0.02816244, 0.15939426, ..., 0.1943017 , + -0.09701379, -0.03432931], + [ 0.10851432, -0.19066858, 0.04841118, ..., -0.01213947, + -0.04722768, 0.09524001]], + + [[ 0.2550473 , -0.04364953, -0.47635067, ..., 0.23234586, + -0.365624 , -0.01943434], + [-0.2977616 , 0.00419303, -0.8656642 , ..., 0.30416554, + -0.04954646, -0.1717201 ], + [ 0.11049867, -0.17219873, -0.6994334 , ..., 0.08241771, + 0.03823497, 0.0737647 ]], + + [[ 0.32358578, 0.17112552, 0.14320694, ..., -0.16366176, + -0.2593694 , 0.07770303], + [-0.1454706 , 0.12558867, -0.0796358 , ..., -0.1381298 , + -0.02053097, 0.09441353], + [ 0.05893492, -0.00259268, 0.02598703, ..., -0.20935844, + 0.1816749 , 0.30294153]]], + + + [[[ 0.14045075, -0.10724663, 0.1179781 , ..., 0.03702212, + 0.30327222, 0.0930538 ], + [-0.30190122, 0.11180247, -0.158114 , ..., 0.0602173 , + 0.28998795, -0.29604238], + [ 0.15810204, -0.02881167, -0.5261955 , ..., 0.16892877, + -0.02681437, -0.18120374]], + + [[ 0.1801053 , -0.17351657, 0.55432343, ..., -0.04310204, + 0.20605011, 0.15464272], + [-0.35629642, 0.1292508 , 0.4196015 , ..., -0.00466895, + 0.40400398, -0.22713828], + [ 0.3233468 , -0.09626149, 0.11579885, ..., 0.07394816, + -0.05250774, 0.03030762]], + + [[ 0.00362034, -0.27520794, -0.2702903 , ..., -0.0462058 , + 0.11103056, -0.17394325], + [-0.5339498 , 0.0243352 , -0.31235537, ..., 0.11845686, + 0.4861063 , -0.42657182], + [ 0.18858455, -0.18266024, -0.27887428, ..., 0.160679 , + 0.1009312 , -0.00895783]], + + [[ 0.23181222, -0.16900635, 0.17172067, ..., -0.2783022 , + 0.0531309 , -0.14735658], + [-0.17012398, 0.0160596 , 0.26340246, ..., -0.15997212, + 0.34474295, -0.18495837], + [ 0.25120145, -0.10030471, 0.4254995 , ..., -0.14183548, + 0.27963412, 0.15028709]]], + + + [[[ 0.04808923, -0.0793573 , -0.05729656, ..., -0.06481453, + -0.08433866, 0.25621825], + [-0.21680164, 0.12764114, -0.13568583, ..., -0.07495075, + -0.22111697, -0.15526403], + [ 0.3291076 , -0.05329537, -0.51923084, ..., 0.05456135, + -0.5494265 , -0.08629624]], + + [[ 0.00642845, -0.02851976, 0.48139414, ..., -0.3939782 , + -0.02460362, 0.4114907 ], + [-0.33011124, 0.21700603, 0.6363917 , ..., -0.34058657, + 0.07507597, -0.00235403], + [ 0.4467702 , -0.08816491, 0.338514 , ..., -0.20894945, + -0.44596782, 0.22598238]], + + [[-0.12848432, -0.17561927, -0.29103017, ..., -0.3294456 , + 0.0409772 , 0.05005151], + [-0.47656342, 0.09599458, 0.01909353, ..., -0.15514864, + 0.3270731 , -0.2028388 ], + [ 0.28001264, -0.17020059, 0.09817386, ..., -0.0626152 , + -0.07981837, 0.1589364 ]], + + [[ 0.18718387, -0.22170076, -0.12480265, ..., -0.2592998 , + -0.3301834 , 0.06356978], + [-0.05257114, -0.01859191, 0.32042652, ..., -0.23058891, + -0.06912984, 0.0211558 ], + [ 0.38780808, -0.15261608, 0.6172939 , ..., -0.2062019 , + -0.08611296, 0.31126183]]], + + + [[[-0.17952503, -0.05265503, -0.25038588, ..., 0.2504428 , + -0.02090847, 0.28380206], + [-0.2667295 , 0.07449401, -0.06040205, ..., 0.07079565, + -0.10550356, 0.09134477], + [ 0.09433349, -0.00351469, -0.22650217, ..., 0.02396235, + -0.11385609, -0.00377663]], + + [[-0.0575775 , 0.03152074, -0.16277562, ..., 0.3681371 , + 0.01227211, 0.08380491], + [-0.17844254, 0.10643482, 0.21974468, ..., 0.26805398, + 0.11603563, -0.11773686], + [ 0.38371345, -0.10176778, 0.20368902, ..., 0.1509468 , + -0.02080609, -0.06920536]], + + [[-0.18148401, -0.06984355, -0.85413545, ..., 0.47789526, + 0.04577804, -0.23575552], + [-0.3161512 , 0.02919929, -0.4008322 , ..., 0.44068208, + 0.26976907, -0.33418196], + [ 0.20420933, -0.1428987 , -0.11891793, ..., 0.30973166, + 0.22936389, -0.15018184]], + + [[-0.08356225, -0.13858244, -0.32757443, ..., 0.142595 , + -0.351965 , -0.0224065 ], + [-0.17430148, -0.03812229, 0.26883033, ..., -0.06516638, + -0.17817669, 0.04958347], + [ 0.09899019, -0.05975993, 0.6693031 , ..., -0.12235346, + 0.04438766, 0.1941795 ]]]], dtype=float32), + }, + }, + position_embeddings: Array([[[ 0.033918 , -0.0929377 , -0.10530281, ..., 0.28548247, + 0.08699364, 0.3533416 ], + [-1.1857578 , -0.05044684, 0.04907949, ..., 0.22999725, + 0.11194737, -0.5982713 ], + [-0.8083621 , -0.45285803, -0.35289305, ..., 0.1278877 , + -0.05470289, -0.3990702 ], + ..., + [ 0.01943148, -0.21790016, 0.20034477, ..., 0.13447031, + -0.12896743, 0.67419606], + [-0.5797325 , 0.09393977, -0.10601121, ..., 0.5656874 , + -0.4588539 , 0.3854702 ], + [-0.5565649 , 0.41388488, 0.21124075, ..., 0.8872757 , + -0.52099276, -0.30215293]]], dtype=float32), + }, + encoder: { + layer: { + 0: { + attention: { + attention: { + key: { + bias: Array([-0.03415372, 0.45620176, -0.6126655 , -0.51253253, 0.4352034 , + 0.649004 , 0.68406653, -0.46161497, -0.42964068, 0.8848562 , + 0.61135024, -0.67950135, -0.08493523, -0.29988992, 0.17360933, + 0.6337579 , 0.08232927, -0.59048873, -0.31515497, 0.09602205, + 0.14608672, -0.06844742, 0.5664715 , -0.6412251 , 0.44181392, + -0.40813568, 0.6399837 , 0.41530073, -0.66080755, -0.01420364, + -0.19360113, 0.14690916, -0.9573473 , -0.3561616 , -0.18794109, + 0.5325726 , -0.5846602 , -0.73184437, -0.35172847, -0.27912554, + 0.33543518, -0.21184462, 0.35979107, -0.12768705, 0.3412426 , + 0.35638675, 0.02133832, -0.08022975, -0.15245748, -0.20536591, + -0.11586879, -1.1425315 , -0.02779106, 0.32926312, -0.26967174, + 0.1869237 , -0.04547051, -1.0505885 , 0.9415589 , -0.5067646 , + -0.4434777 , 0.31517163, -0.48155847, 0.53517747, -0.40972438, + 0.16880448, 0.17367116, 0.29588723, 0.05554008, -0.48776305, + -0.31501368, -0.10502674, 0.13945673, 0.55460685, 0.51986885, + -0.39894032, -0.2494641 , 0.05012343, -0.04834368, 0.29485628, + 0.19806208, 0.3374114 , 0.11331975, 0.36502162, -0.01847783, + -0.39011198, 0.3790581 , 0.02779981, 0.651827 , 0.46813282, + 0.29877353, -0.43409863, 0.15295495, -0.35798907, -0.43806 , + -0.26218414, -0.45440933, 0.11940431, -0.4274145 , 0.18455578, + 0.00287167, 0.33210322, 0.32568368, -0.33072832, 0.08617178, + -0.24881381, -0.07380535, 0.23978767, -0.12747152, -0.7031561 , + 0.10753731, 0.27640837, 0.0658379 , -0.00729593, 0.15286936, + 0.14529614, 0.20602514, -0.18859431, 0.08264172, 0.09762905, + -0.1777643 , 0.21521233, 0.68952733, 0.35572144, -0.25402895, + -0.12086129, 0.33329782, -0.27421075], dtype=float32), + kernel: Array([[-0.13037314, -0.3519865 , -0.7311013 , ..., -0.14157563, + -0.098704 , 0.08183759], + [ 0.36583012, 0.46603498, 0.07223602, ..., 0.3872796 , + -0.01522054, 0.1815025 ], + [ 0.4379942 , -0.2856487 , -0.14937577, ..., 0.2555367 , + 0.143062 , -0.10283173], + ..., + [-0.08320192, -0.2315665 , 0.183584 , ..., -0.08580804, + -0.13162506, 0.07868645], + [ 0.42390886, -0.22044583, 0.4916073 , ..., 0.57441944, + -0.04296773, 0.02346938], + [ 0.09392057, -0.64890575, -0.21604928, ..., 0.09232247, + 0.13443615, 0.11167435]], dtype=float32), + }, + query: { + bias: Array([-0.06905509, -0.44198722, 0.30581513, 0.14314957, -0.15857255, + 0.15910107, -0.14257118, 0.11895534, -0.13231936, -0.28292307, + -0.1729494 , 0.1297792 , 0.29567292, 0.1645278 , -0.07054784, + -0.1081412 , 0.2331662 , -0.02096415, 0.00768806, -0.13316594, + -0.25273684, -0.09793813, 0.04979343, 0.23559354, -0.19012289, + 0.08803113, -0.03696692, -0.29770622, 0.39657098, 0.1723672 , + 0.12946679, 0.12689376, 0.3535053 , -0.02396139, 0.01831555, + -0.45482886, 0.0035749 , 0.0489781 , -0.19142908, 0.04616706, + -0.10055473, 0.00193944, -0.0123636 , -0.1297169 , 0.06566181, + -0.12390063, 0.01915232, -0.10427503, -0.26933163, 0.01553594, + -0.02893843, 0.49531803, 0.30729276, 0.01837071, 0.172419 , + -0.05943993, -0.29665995, 0.46334326, -0.13110274, 0.49487928, + 0.21740474, -0.25742924, 0.03184005, -0.16928338, 0.04680282, + -0.10523521, 0.07657482, 0.147745 , -0.09269993, 0.01065696, + 0.05368982, -0.18941729, 0.10891443, 0.20821781, 0.15085766, + 0.51667964, -0.15096618, 0.34021842, -0.03054576, 0.26622248, + -0.36788234, 0.29284194, -0.18616207, 0.09047277, 0.3589835 , + -0.09647398, -0.10208427, 0.09295312, 0.03126433, -0.27378327, + 0.08303533, -0.17912042, -0.06608632, -0.05499709, -0.09980894, + 0.15720649, 0.37134722, 0.39746282, 0.172588 , 0.04140972, + 0.58430547, -0.250469 , -0.3090034 , 0.35163316, -0.09326122, + 0.18299101, -0.02789885, -0.16457352, 0.21470276, 0.20549582, + -0.27671963, -0.4622744 , 0.19656613, 0.31474093, -0.0242331 , + -0.01637263, 0.08560522, 0.40495184, -0.14052592, -0.24214344, + -0.27252468, 0.15204628, -0.19363014, -0.13963734, -0.02372985, + -0.13916686, 0.21864863, 0.13198575], dtype=float32), + kernel: Array([[-0.02407378, -0.08052919, -0.78489935, ..., 0.0214601 , + -0.21603775, 0.07814599], + [ 0.1441562 , 0.39637634, -0.01425371, ..., 0.44336745, + 0.02438617, 0.04486706], + [ 0.54582417, -0.38129556, -0.09409017, ..., 0.07300002, + 0.08039478, -0.01730129], + ..., + [-0.40244198, 0.23577765, -0.09792052, ..., 0.06759664, + 0.0144011 , -0.21112448], + [ 0.52981406, 0.07420167, 0.320484 , ..., 0.3295426 , + -0.08885895, -0.11461516], + [ 0.2322218 , -0.26398233, -0.4191451 , ..., 0.24738975, + -0.05001807, 0.01805569]], dtype=float32), + }, + value: { + bias: Array([-2.94787940e-02, -2.72796545e-02, 6.14140369e-02, -2.86302455e-02, + 7.17062429e-02, -3.11819911e-02, -4.36154753e-02, 3.26400585e-02, + -2.82624215e-02, 1.26899080e-02, -7.74416178e-02, 4.88068024e-03, + 2.13063625e-03, -7.00014010e-02, 6.05225153e-02, -4.30300795e-02, + 1.79692768e-02, -2.71751042e-02, 3.53089422e-02, -1.95429940e-02, + -1.21456282e-02, -2.05201493e-03, -3.62878200e-05, -1.08139282e-02, + 7.40849450e-02, 1.50353080e-02, -7.06956685e-02, 8.49624351e-02, + -7.26162344e-02, -5.55856004e-02, -7.41069913e-02, 7.11530168e-03, + -2.33412739e-02, 6.24122694e-02, 1.00617625e-01, 4.78886962e-02, + 6.93868250e-02, 8.44263211e-02, 6.89995065e-02, -1.49511993e-02, + 6.76077977e-02, 2.99593173e-02, -6.50391206e-02, -3.10420450e-02, + -4.87265997e-02, -7.07542598e-02, 3.83699648e-02, -1.07949436e-01, + 3.88861708e-02, 1.07101694e-01, -7.05710873e-02, 6.53532194e-03, + 5.76992594e-02, 8.30659941e-02, -1.04475126e-01, 8.40004310e-02, + -9.89107862e-02, 8.94425064e-02, -3.12575921e-02, -5.92782833e-02, + -8.53339434e-02, -5.87267503e-02, 3.41176279e-02, -6.32282570e-02, + -4.08555977e-02, -7.24948272e-02, 4.16365713e-02, -2.20883206e-01, + -4.90566827e-02, -1.89475138e-02, -1.50290072e-01, 2.09986553e-01, + 3.20502929e-02, 8.28777160e-03, 6.13699071e-02, -2.50448077e-03, + -2.15128988e-01, -3.07804886e-02, -1.82573318e-01, 9.74187329e-02, + 8.86107534e-02, -6.78683072e-02, -6.89225420e-02, -3.00914403e-02, + -1.05735838e-01, -2.17407808e-01, 1.63544074e-01, 1.00356489e-01, + -6.03550719e-03, 3.53857167e-02, -7.28843063e-02, 1.15036063e-01, + -1.29502550e-01, 5.06461645e-03, 1.20808132e-01, 1.27421441e-02, + -3.18243355e-02, -5.57740144e-02, -6.90745190e-02, -1.28188869e-02, + 1.20238643e-02, -4.42795120e-02, -2.79449914e-02, -3.33387293e-02, + 7.23095536e-02, -2.58487016e-02, -1.62459761e-02, 1.07147522e-01, + -6.11629374e-02, 7.07705468e-02, -5.56649491e-02, 3.85964476e-02, + 9.34836396e-04, -1.59531720e-02, -7.20297545e-02, 8.05353299e-02, + 1.81769580e-02, -3.95291820e-02, -7.38437474e-02, -9.54268277e-02, + 1.34773731e-01, 9.29158255e-02, -5.70693836e-02, 2.48423647e-02, + -8.44348520e-02, -7.13706389e-02, -7.08770528e-02, 7.93647394e-02], dtype=float32), + kernel: Array([[-0.19559653, 0.10921917, -0.24855891, ..., -0.3095954 , + -0.11501747, 0.3495971 ], + [ 0.2759333 , -0.0968355 , 0.25864232, ..., -0.07446828, + -0.01653024, 0.04277551], + [-0.53325766, -0.17772451, -0.19101648, ..., -0.09837775, + 0.07115 , 0.39424804], + ..., + [ 0.36891153, -0.11530672, 0.24386573, ..., 0.05493925, + -0.2263004 , -0.10157143], + [ 0.06181217, 0.21771231, -0.52629954, ..., -0.0283559 , + -0.11670491, -0.19776788], + [-0.20297265, -0.24884045, 0.21022551, ..., 0.13422379, + 0.22147001, 0.18063086]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-0.08976492, 0.35623267, -0.28853896, 0.1584903 , -0.04931528, + 0.41930285, -0.12198956, -0.1676379 , -0.07629202, 0.44505563, + -0.13279814, -0.2606873 , -0.38957864, 0.2600788 , -0.15301487, + -0.7227959 , 0.06558075, -0.02327237, -0.01965788, 0.3900318 , + 0.14463983, 0.33909434, -0.02560525, 0.25102413, 0.26336563, + 0.07151037, -0.01791747, -0.32692116, -0.29931375, 0.26871207, + -0.45228213, -0.07384168, -0.2372598 , 0.16236512, 0.20462725, + -0.04373658, 0.18107706, 0.5327047 , 0.11542407, -0.20473395, + -0.29505363, -0.1296326 , -0.50376815, 0.23728877, 0.16542588, + 0.20167176, 0.05926127, 0.1154462 , -0.14579555, -0.08487356, + 0.07978057, -0.27115968, -0.5287936 , 0.07883674, -0.07367402, + -0.04930678, -0.5660592 , 0.14252122, -0.05342535, -0.03159878, + 0.18108739, -0.06234799, 0.13124788, 0.00128543, -0.24720064, + 0.13804607, -0.5026876 , 0.53043497, 0.07013435, 0.22524473, + 0.17431633, 0.22486666, 0.01429994, -0.02526889, 0.65786517, + -0.12362218, -0.124548 , 0.02234903, -0.16374543, -0.03648356, + -0.79906535, -0.21115243, -0.20314318, 0.24117418, -0.42399853, + 0.37275848, 0.11345251, -0.07611627, 0.09416302, -0.16610427, + 0.22861136, -0.22463125, 0.21705595, 0.09253941, -0.2043051 , + -0.06828043, -0.27955458, -0.00722787, -0.13780643, -0.44847608, + 0.18490234, 0.09589677, 0.2274974 , -0.04555204, -0.13882257, + -0.17222178, -0.24243575, 0.09366386, 0.16359329, 0.15961671, + 0.08576287, 0.09707799, 0.31846833, -0.27176753, 0.12482687, + -0.35437414, -0.04145558, -0.4648349 , 0.1269856 , -0.09939937, + 0.00209725, 0.00080232, 0.20968738, -0.18917999, -0.12978964, + 0.13607804, 0.00752292, -0.00256601], dtype=float32), + kernel: Array([[ 0.01726927, 0.12866902, -0.1439863 , ..., -0.26570368, + -0.0510007 , 0.47683123], + [ 0.15968257, -0.1014777 , -0.26105717, ..., 0.30918548, + 0.48335168, 0.01941068], + [-0.05697384, 0.05050297, -0.3797851 , ..., -0.70562077, + 0.11496508, 0.02426186], + ..., + [-0.37025422, -0.41289684, -0.09372032, ..., 0.3905355 , + -0.44608384, -0.10514601], + [ 0.03965069, -0.21772908, -0.06686966, ..., 0.5678079 , + -0.5982024 , 0.03439448], + [ 0.31467485, 0.03231593, 0.6761289 , ..., 0.3369738 , + -0.16577064, 0.09921929]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([ 7.65252188e-02, 6.29107952e-01, 8.02198052e-02, 8.70596059e-03, + 4.38350849e-02, -1.07495952e-02, -5.31466417e-02, -1.02623366e-01, + 2.12361347e-02, 6.80099875e-02, -2.28481810e-03, 8.53950828e-02, + -5.73644787e-02, 2.06265952e-02, -7.53715783e-02, -1.43298656e-01, + -2.74074376e-02, -1.09807447e-01, -6.17579147e-02, -7.90846869e-02, + -9.07771438e-02, 8.38840082e-02, -1.48849720e-02, -2.24489588e-02, + -9.19756591e-02, -1.33614585e-01, 8.40165392e-02, -2.04390988e-01, + 7.47799221e-03, 1.38857782e-01, -5.69530539e-02, 9.17964336e-03, + -1.55208111e-01, 2.73925383e-02, -1.95245608e-03, 5.41600175e-02, + -2.93773953e-02, 3.53689879e-01, -1.23494409e-01, -1.56066343e-01, + 8.71012285e-02, 2.55195685e-02, -3.10201854e-01, -8.39867163e-03, + -3.08605190e-02, -5.28353490e-02, -1.40786827e-01, -5.47004975e-02, + -1.13185709e-02, 1.86418276e-02, -7.84755684e-03, -4.15783050e-03, + -7.23032653e-02, -1.61183357e-01, 1.90100819e-02, -4.55883006e-03, + -2.70947274e-02, 1.69784650e-02, -2.28463719e-03, 4.21662033e-02, + -9.98298172e-03, 7.62146562e-02, -1.09199071e-02, 7.20275193e-02, + 1.14038222e-01, 2.85849470e-04, -2.99424119e-02, 1.37919337e-01, + -1.45379841e-01, 6.62748562e-03, 2.23827194e-02, -5.66348387e-03, + 5.01772054e-02, 1.41351158e-03, 5.96981421e-02, -2.08293200e-01, + -5.95040731e-02, 5.53276669e-03, -8.04201216e-02, -1.67969614e-01, + 3.61396512e-03, 5.54506434e-03, -2.92952433e-02, 4.65981402e-02, + 1.59924058e-03, -2.43443865e-02, -1.17824182e-01, -3.83222140e-02, + -3.87748890e-02, -6.59047961e-02, -8.08075536e-04, -3.35978210e-01, + -4.24195863e-02, 3.55881304e-01, 3.02827614e-03, -5.55533350e-01, + 1.54737878e-04, -1.47938222e-01, -2.13265121e-02, -1.14879206e-01, + -3.36717785e-04, 5.57862930e-02, 1.95077602e-02, -8.35356787e-02, + 2.89721284e-02, -1.48645073e-01, -2.94460828e-04, 8.76041055e-02, + -1.90158889e-01, 1.32130921e-01, 1.11547068e-01, -8.59021768e-02, + -4.04770300e-03, -2.04068184e-01, 2.25875482e-01, -7.98009858e-02, + 1.03451028e-01, -6.95376145e-03, -8.43686834e-02, -1.78507604e-02, + -9.83250737e-02, -7.10767433e-02, -6.19554799e-03, -5.51361814e-02, + -8.06146394e-03, 4.24989313e-01, -2.02743262e-02, -7.10928217e-02], dtype=float32), + scale: Array([ 4.92835701e-01, 7.25533962e-01, 1.97780564e-01, 7.19367981e-01, + 3.85366410e-01, 1.70830980e-01, 4.03196029e-02, 3.59184593e-01, + 1.64954796e-01, 5.99162340e-01, 1.94611892e-01, 4.47059900e-01, + 1.33249581e-01, 1.17341705e-01, 2.57417858e-01, 2.13419840e-01, + 1.17401130e-01, 1.44011736e-01, 4.55426246e-01, 3.37985694e-01, + 3.68331194e-01, 1.58645034e-01, 4.11303312e-01, 2.80974478e-01, + 2.08782077e-01, 3.16983044e-01, 2.54319280e-01, 5.07932603e-01, + 6.26261353e-01, 3.31194192e-01, 5.77102304e-01, 5.91307692e-03, + 2.17876673e-01, 1.83340341e-01, 7.28062959e-03, 1.43880785e-01, + 3.79934996e-01, 8.64668906e-01, 5.03845096e-01, 3.30553919e-01, + 3.66227090e-01, 1.23913087e-01, 4.04930204e-01, -5.25120134e-03, + 2.16952696e-01, 2.82419056e-01, 2.41874039e-01, 2.18594924e-01, + 2.72880584e-01, 1.63890943e-01, 5.69228344e-02, 1.57258492e-02, + 3.11221033e-01, 3.75806093e-01, 3.45861703e-01, 3.33030730e-01, + 1.86815023e-01, 2.87897736e-01, 1.79882497e-02, 3.26225400e-01, + 3.56188752e-02, 3.17295760e-01, 5.93220115e-01, 3.10718954e-01, + 1.92407712e-01, 9.49200336e-03, 2.15880141e-01, 4.15735185e-01, + 3.10675055e-01, 5.28817117e-01, 1.04173638e-01, 1.57576382e-01, + 3.02691996e-01, 3.79474498e-02, 7.07599940e-03, 5.08240998e-01, + 3.50080490e-01, 3.61189306e-01, 1.87464565e-01, 4.86818105e-01, + 3.34662758e-03, -1.36662647e-03, 1.20725386e-01, 5.04414022e-01, + 3.98086966e-04, 9.34311152e-02, 1.42500699e-01, 3.84391755e-01, + 1.07679166e-01, 3.01436037e-01, -1.05553423e-03, 4.41559106e-01, + 1.26928464e-01, 4.32251424e-01, -6.94631459e-03, 5.02396822e-01, + -3.20078339e-04, 1.94142312e-01, 8.77228081e-02, 1.18769050e-01, + 4.47774172e-01, 3.87618184e-01, 4.12855119e-01, 3.15822661e-01, + 4.97932345e-01, 6.33932471e-01, 6.15001982e-03, 2.09667221e-01, + 2.26229250e-01, 1.05634868e+00, 5.23372293e-01, 4.29101616e-01, + 2.27578685e-01, 6.58651769e-01, 2.97994912e-01, 1.07822992e-01, + 4.83628422e-01, 4.72095191e-01, 3.23185503e-01, 1.00680351e-01, + 4.04555202e-01, 1.64712802e-01, 6.41667992e-02, 1.81472510e-01, + 5.06515145e-01, 4.70199496e-01, 2.70896405e-01, 4.36434001e-01], dtype=float32), + }, + layernorm_before: { + bias: Array([-0.05890432, 0.07051883, 0.09445145, -0.07912993, -0.14302658, + -0.01429262, 0.03426607, 0.00171352, 0.11276561, -0.05662137, + 0.13313785, -0.04794319, 0.00357568, 0.01197018, -0.07405796, + -0.06490814, -0.0245945 , 0.14022568, -0.05081064, -0.12518148, + 0.00789185, -0.24005446, 0.02322139, -0.17626956, -0.03338344, + -0.33740795, -0.13774505, -0.07997898, 0.04349966, -0.01520896, + -0.39623648, -0.150468 , 0.11885046, -0.01181321, 0.00808169, + -0.20219389, 0.0043777 , -0.08942726, -0.16421181, -0.18974544, + -0.07742172, -0.01755824, -0.02627599, -0.4378331 , -0.19054742, + 0.2720697 , 0.25448796, 0.01773962, -0.01873611, 0.0169231 , + 0.03406947, 0.0349858 , 0.09072868, -0.05627641, 0.14739008, + 0.01965971, 0.09791738, 0.03601414, -0.28183195, 0.15014632, + 0.22542529, 0.3072184 , -0.09403451, -0.08351854, 0.25167593, + 0.04495538, -0.11564947, -0.00397131, 0.10662062, 0.05204451, + 0.01583207, 0.01414629, 0.23943776, 0.06893958, -0.08743864, + -0.06693614, -0.07541402, -0.19011301, -0.17251883, 0.01803088, + -0.09193094, 0.04894969, -0.00502513, -0.09108073, -0.1768744 , + -0.01100582, 0.10178719, 0.22854435, 0.04437195, 0.00859537, + -0.03013818, 0.26522186, -0.0835686 , 0.0909194 , 0.08486014, + -0.39586094, 0.03491055, 0.02060359, -0.14571564, -0.18893251, + 0.04928007, 0.05034845, -0.15555385, -0.15786135, -0.56710386, + 0.01834173, -0.04088113, -0.0075858 , 0.04828192, -0.1228609 , + -0.00563954, -0.3955512 , -0.06544469, 0.09136705, -0.10234332, + 0.09786685, 0.00981475, -0.05797027, 0.08497368, 0.18863663, + -0.00070434, 0.03551389, -0.04786862, -0.03968854, -0.02606042, + -0.2953443 , -0.03344315, -0.2232286 ], dtype=float32), + scale: Array([0.45335642, 0.28725365, 0.25888288, 0.51014036, 0.25393826, + 0.14649442, 0.45231742, 0.07753127, 0.3569898 , 0.67538893, + 0.2085993 , 0.17448373, 0.01005436, 0.07782084, 0.26946452, + 0.14676209, 0.22900493, 0.29659525, 0.13212137, 0.1925922 , + 0.04533805, 0.26069018, 0.15647988, 0.20724292, 0.08556735, + 0.296176 , 0.20237958, 0.6578139 , 0.11472861, 0.3509395 , + 0.38030067, 0.19699657, 0.30957213, 0.1079342 , 0.12253352, + 0.54774934, 0.06868377, 0.49437425, 0.34611228, 0.33752495, + 0.24485146, 0.11372902, 0.13047746, 0.40855378, 0.49101192, + 0.5738504 , 0.32902393, 0.2079778 , 0.09577054, 0.11233962, + 0.17488489, 0.10529607, 0.13811632, 0.19155163, 0.16904642, + 0.10060844, 0.2304541 , 0.20629783, 0.21513213, 0.1885996 , + 0.2502167 , 0.3126915 , 0.3413607 , 0.1449984 , 0.41228867, + 0.09932974, 0.14989796, 0.13950281, 0.24986736, 0.15920106, + 0.06112578, 0.10076807, 0.25724828, 0.11337321, 0.1927693 , + 0.24177718, 0.14134662, 0.4839469 , 0.20330684, 0.1124863 , + 0.17910473, 0.1306797 , 0.09654108, 0.13909315, 0.19018278, + 0.12305983, 0.3541507 , 0.31824607, 0.13942654, 0.08698463, + 0.08841392, 0.17479059, 0.15011418, 0.45322958, 0.17411204, + 0.29569057, 0.10512718, 0.3171436 , 0.21283199, 0.51158214, + 0.42643854, 0.18322334, 0.3302157 , 0.37901062, 0.44037428, + 0.50349325, 0.1318383 , 0.13000993, 0.14093749, 0.5986485 , + 0.1306848 , 0.36673996, 0.2426187 , 0.19750752, 0.14503735, + 0.22176953, 0.1925638 , 0.15475766, 0.17507745, 0.22627683, + 0.10930963, 0.37122646, 0.12475897, 0.12880948, 0.11971601, + 0.2903781 , 0.1446715 , 0.20750394], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-0.4722571 , -0.43928552, -0.5569621 , -0.3837599 , -0.47573015, + -0.4316669 , -0.762032 , -0.5649565 , -0.31526604, -0.42409667, + -0.35039356, -0.5466656 , -0.08653939, -0.50116634, -0.65188706, + -0.47724712, -0.29790252, 0.32602024, -0.04504937, -0.40849677, + -0.11292349, -0.25168103, -0.45442316, -0.34808317, -0.38236308, + -0.3788954 , -0.19183853, -0.30315536, -0.16824003, -0.4485127 , + -0.30100617, -0.39639288, -0.16540907, -0.5450078 , -0.4645357 , + -0.36592278, -0.5585429 , -0.34257036, -0.19063933, -0.44469953, + -0.42737913, -0.36188826, -0.2907807 , -0.02146477, -0.23154494, + -0.20253362, -0.69757545, -0.39887667, -0.44025323, -0.39642838, + -0.387052 , -0.24472845, -0.09177764, 0.05788963, -0.04718702, + -0.21287683, -0.29563203, -0.3253591 , -0.43122077, -0.1849469 , + -0.5141688 , -0.29721844, -0.47572163, -0.3446094 , -0.26447332, + -0.47591433, -0.19857107, -0.6629566 , -0.36350504, -0.60418856, + -0.4550368 , -0.40432703, -0.35896167, -0.33915696, -0.41339654, + -0.23302272, -0.05349919, -0.47071078, -0.32239798, -0.3542714 , + -0.25527537, -0.3346329 , -0.49991933, -0.22825265, -0.45629248, + 0.01691399, -0.34854415, -0.44188562, -0.03925477, -0.3642511 , + -0.40468022, -0.33773935, -0.49415597, -0.5592706 , -0.2977098 , + -0.32126102, -0.16516496, -0.40793797, -0.8022191 , -0.44280693, + -0.1008477 , -0.22637095, -0.16032524, -0.45757398, -0.4884214 , + -0.48628846, -0.4652781 , -0.51417524, -0.25536695, -0.33155087, + -0.5587205 , -0.37767413, -0.30986476, -0.47982594, -0.51307243, + -0.38967907, -0.24242653, -0.25831246, -0.1458375 , -0.58483136, + -0.3381596 , 0.27927315, -0.7453308 , -0.4533877 , -0.15844217, + -0.45824784, -0.468726 , -0.4619875 , -0.4371935 , 0.00825728, + -0.31616744, -0.4558102 , 0.0993996 , -0.08013049, -0.17536967, + -0.52361816, -0.44187388, -0.25537458, -0.19660369, -0.30101493, + -0.09137805, -0.41055298, -0.46738333, -0.57409215, -0.34877914, + -0.52286893, -0.55365384, -0.33294562, -0.47346646, -0.3105509 , + -0.44117743, -0.3502154 , -0.39879954, -0.3967916 , -0.47747064, + -0.24431925, -0.39359137, -0.22136432, -0.35870054, -0.29925522, + -0.34952432, -0.49339256, -0.42442355, -0.383239 , -0.3228079 , + -0.59250957, -0.3766234 , -0.48536432, -0.49891275, -0.10365842, + -0.21082799, -0.5129332 , -0.39634183, -0.4714169 , -0.3876767 , + -0.2642752 , -0.1780093 , -0.08995035, -0.5045371 , -0.31991875, + -0.3657458 , -0.51370513, -0.5282006 , -0.43932983, -0.44844607, + -0.4636747 , -0.541842 , -0.28016394, -0.5640383 , -0.28513128, + -0.44549647, -0.43519762, -0.42721725, -0.57123166, -0.40650263, + -0.12434424, -0.2684522 , -0.13034356, -0.41602165, -0.2656361 , + -0.63166195, -0.24252081, -0.49084774, -0.66598415, -0.21980134, + -0.20069678, -0.3857888 , -0.5750186 , -0.28702593, -0.42044863, + -0.43332833, -0.05599573, -0.17978549, -0.33969852, -0.42026928, + -0.19741155, -0.2731103 , -0.5647213 , -0.61824954, -0.44372314, + -0.16089505, -0.6141717 , -0.27717924, -0.39058053, -0.30320698, + -0.3017265 , -0.36256546, -0.29453766, -0.7215134 , -0.36851448, + -0.45460737, -0.25023326, -0.24756469, 0.5198105 , -0.49648827, + -0.7883031 , -0.40549907, -0.44536465, -0.34158024, -0.4282217 , + -0.45166638, -0.285461 , -0.5351079 , -0.5246095 , -0.38692436, + -0.4037612 , -0.68792176, 0.09384751, -0.3728323 , -0.54188645, + -0.14435874, -0.13926642, -0.48824954, 0.103251 , -0.5347266 , + -0.12768686, -0.2596208 , -0.33724293, 0.41711676, -0.46974605, + -0.31142974, -0.46029773, -0.38704032, -0.4837437 , -0.4762127 , + -0.48931426, -0.27646464, -0.7709319 , -0.11419559, -0.3580196 , + -0.41691482, -0.523539 , -0.7911689 , -0.14548191, -0.32703182, + -0.22812605, -0.5490665 , -0.27162734, -0.47042787, -0.27687344, + -0.29244688, -0.14437722, -0.42082107, -0.46603575, -0.3757466 , + -0.43942484, -0.35704526, -0.50712955, -0.06310245, -0.51886773, + -0.36505163, -0.29938784, -0.47636163, -0.798845 , -0.39453545, + -0.42888972, -0.34777957, -0.42616373, -0.37013018, -0.2707561 , + -0.22923009, -0.4484783 , 0.0781149 , -0.54366326, -0.31296468, + -0.46913856, -0.26365936, -0.05434812, -0.43664017, -0.262903 , + -0.21987678, -0.29727325, -0.53905886, -0.2881995 , -0.2697021 , + -0.56999284, -0.23401402, -0.25733256, -0.48550016, -0.53915155, + -0.29186913, -0.41333848, -0.11731004, -0.1876421 , -0.38366327, + -0.69185966, -0.54369533, -0.13784738, -0.45933396, -0.47800148, + -0.42821413, -0.37634867, -0.42712605, -0.06954894, -0.445592 , + -0.40655133, -0.3772172 , -0.7594351 , -0.5565112 , -0.25855637, + -0.13048875, -0.3732955 , -0.30561906, -0.5940574 , -0.41309616, + -0.314704 , -0.3010439 , -0.5837942 , 0.40323326, -0.24861278, + -0.4630232 , -0.538137 , -0.37435848, -0.46322653, -0.48135036, + -0.26899594, -0.17752594, -0.46429896, -0.62554 , -0.4089612 , + -0.34856302, -0.07396294, -0.4918398 , -0.44181666, -0.21933144, + -0.41612074, -0.46780694, -0.25013933, -0.3507257 , -0.31338322, + -0.49207342, -0.573551 , -0.25965518, -0.27235296, -0.46362886, + -0.572643 , -0.36132613, -0.5891415 , -0.43028024, -0.58089757, + -0.5628406 , -0.23266539, -0.48305756, -0.4171379 , -0.30668163, + -0.800837 , 0.16001745, -0.32928717, -0.7041144 , -0.28114614, + -0.1339794 , -0.1539137 , -0.48452973, -0.2916893 , -0.31748015, + -0.41399214, -0.5892567 , -0.6904999 , -0.39504105, -0.05394559, + -0.4220559 , -0.08093858, -0.55065817, -0.42636526, -0.36694276, + -0.6491974 , -0.16998939, -0.6064266 , -0.49446252, -0.25082645, + -0.3862861 , -0.5665641 , -0.12244188, -0.59443516, -0.43923312, + -0.03984782, -0.4539737 , -0.2804616 , -0.28886226, -0.41679534, + -0.7231176 , -0.44628784, -0.17072082, -0.55711204, -0.33631635, + -0.3513167 , -0.44696072, -0.09838888, -0.5331628 , -0.3244455 , + -0.22034623, -0.39987078, -0.52567166, -0.18180668, -0.1563768 , + -0.15176131, -0.34147656, -0.39736265, -0.22072102, -0.16304936, + -0.4934699 , -0.35189024, -0.38949656, 0.13233016, -0.2961783 , + -0.03038953, -0.19008778, -0.43870625, -0.14041252, -0.46124077, + -0.6101991 , -0.78746367, -0.41358054, -0.0965796 , -0.17454897, + -0.46205354, -0.5193903 , -0.43612164, -0.338933 , -0.2509465 , + -0.25011972, -0.32999432, -0.5068436 , -0.40737084, -0.59392023, + -0.08752135, -0.6887468 , -0.02706476, -0.34656492, -0.45208535, + 0.12221402, -0.07045163, -0.6397726 , -0.32004473, -0.2506177 , + -0.41133684, -0.20769219, -0.3077098 , -0.48824766, -0.59370387, + -0.3106976 , -0.09882683, -0.23733851, -0.5397758 , -0.384874 , + -0.48532894, -0.35305282, -0.44897798, -0.57191414, -0.23201816, + -0.40165707, -0.45165244, -0.30947968, -0.57236695, -0.40093747, + -0.610941 , -0.4310733 , 0.00399699, -0.17124976, -0.33956885, + -0.59238917, -0.4834356 , -0.23509572, -0.3706481 , -0.46995562, + -0.30272713, -0.42109653, -0.16173697, -0.3515516 , -0.46951428, + -0.24467751, -0.37127313], dtype=float32), + kernel: Array([[-0.16517928, -0.06662413, -0.32296306, ..., 0.17395326, + -0.00140762, -0.21535204], + [-0.2497043 , -0.32800463, -0.54542446, ..., -0.37079784, + -0.4540639 , -0.2762963 ], + [ 0.20276228, -0.2351077 , -0.4339836 , ..., -0.07039763, + 0.09680251, -0.00943627], + ..., + [-0.0865024 , -0.07723805, -1.0022317 , ..., -0.22161727, + 0.2631157 , -0.2520337 ], + [-0.06232752, 0.41687462, 0.43454677, ..., 0.17987384, + -0.22243185, 0.2732002 ], + [-0.1515276 , 0.08726911, -0.16446725, ..., -0.1428107 , + -0.12378144, -0.24571244]], dtype=float32), + }, + output: { + bias: Array([-2.27567896e-01, 3.73193085e-01, 1.41862273e-01, 8.60335305e-02, + 6.99756369e-02, -1.08260542e-01, -7.94294663e-03, 1.12044938e-01, + -1.97965074e-02, 1.03512242e-01, 1.25780692e-02, 2.14173347e-02, + -1.95132956e-01, -1.66688226e-02, -5.57698980e-02, 2.21790038e-02, + 8.23988672e-03, -1.56303987e-01, 2.77717095e-02, -6.28630817e-02, + -1.10570341e-01, 1.39788479e-01, -9.08712447e-02, -7.11519644e-03, + -1.21463031e-01, 1.24595746e-01, 6.99105812e-03, 6.04419075e-02, + 1.16523303e-01, -3.08515457e-03, 1.71046168e-01, -4.62258421e-02, + 1.37897981e-02, -9.23820660e-02, 1.22136414e-01, -5.12558557e-02, + -4.16860264e-03, 5.67508861e-02, -8.22613947e-04, -1.72597811e-01, + -1.73500534e-02, -1.02644004e-01, -7.16797352e-01, 6.50356859e-02, + -9.02397409e-02, -7.97297060e-02, -3.20093259e-02, 7.92354718e-02, + -5.91707882e-03, -5.96906338e-03, 2.83249887e-03, 4.21253145e-02, + -1.67518064e-01, -2.99427286e-02, -5.09011634e-02, -3.20116989e-02, + -1.33028731e-01, 2.27317363e-02, 1.29082099e-01, 1.18409574e-01, + 1.57172829e-02, -2.19231278e-01, -2.51264381e-03, -4.79261391e-02, + 4.22019213e-02, -7.78303593e-02, -7.12049752e-02, -8.38479102e-02, + -3.88463922e-02, -3.15606929e-02, -6.03214689e-02, 4.97607961e-02, + -3.79606709e-02, 7.08331587e-04, 9.65211987e-01, 1.56220675e-01, + -2.14344375e-02, -2.15863492e-02, -1.17833465e-01, 1.19182810e-01, + -2.71626860e-01, 1.97731983e-03, 3.98320071e-02, 2.02119559e-01, + 6.09681681e-02, -3.32502718e-03, 1.39364108e-01, 5.88153526e-02, + -4.45023738e-02, -6.55459762e-02, 4.42181528e-02, -8.78035650e-02, + -4.80993986e-02, -1.09087467e-01, -9.66937765e-02, -1.30174085e-01, + 2.12008785e-03, 2.63354164e-02, -9.37493294e-02, 7.85378218e-02, + 1.73281968e-01, -7.77669847e-02, 1.33560196e-01, -1.09741293e-01, + 1.85674503e-01, 1.03994340e-01, -1.89410765e-02, -1.18796170e-01, + -1.62400916e-01, 2.00676285e-02, -1.06416032e-01, 1.79987997e-02, + 1.19247243e-01, -2.42678579e-02, 1.64977908e-01, -7.41590792e-03, + -1.76154330e-01, 2.93482076e-02, -6.35338202e-02, 2.07892321e-02, + 8.92439187e-02, 2.59270482e-02, 2.71555986e-02, -7.91539550e-02, + 9.35807526e-02, 5.84603362e-02, 1.93760265e-02, 1.72683045e-01], dtype=float32), + kernel: Array([[ 6.6307461e-01, -2.6014799e-01, -8.3183780e-02, ..., + -2.6699215e-01, -1.0033758e-01, -2.1524315e-01], + [-1.4059058e-01, -4.3626988e-01, -8.1478050e-03, ..., + -3.1092042e-01, 3.1286913e-01, -3.7179798e-02], + [ 7.0444113e-01, -1.0556934e-05, -2.0829076e-02, ..., + -9.0309918e-01, 3.2471472e-01, 5.0782524e-02], + ..., + [-7.5270250e-02, 3.7796192e-02, -5.1213449e-01, ..., + 3.4989411e-01, -1.0396782e-01, -5.7098109e-01], + [ 9.8383605e-02, -1.5289782e-01, 3.3700535e-01, ..., + 5.9063655e-01, -9.2052268e-03, -1.1524218e-01], + [ 1.1272128e-01, -6.6285774e-02, -2.5734550e-01, ..., + -1.2502365e-01, 9.1840588e-02, 2.3722002e-01]], dtype=float32), + }, + }, + }, + }, + 1: { + attention: { + attention: { + key: { + bias: Array([ 0.08042084, 0.17798816, 0.05449522, 0.09923156, 0.02864883, + -0.06706466, 0.12987332, 0.05871709, 0.05153077, 0.05436414, + 0.05116336, -0.14068684, -0.11056656, -0.00991776, 0.00717054, + 0.0770971 , 0.11563851, 0.02029003, 0.02723394, -0.12056217, + -0.06265891, -0.05970899, 0.14154014, -0.03163632, 0.0481686 , + 0.02711394, 0.04412205, 0.02307438, -0.05139432, 0.03842278, + -0.07357287, -0.08965195, -0.08623472, 0.03718801, 0.84410733, + 0.4109549 , 0.3907228 , -0.44974568, -0.3517945 , 0.03430397, + -0.07853877, -0.30442524, 0.5081588 , 0.15563154, 0.17845473, + 0.14964798, 0.01935043, 0.78861445, 0.13657925, 0.9571119 , + -0.15077294, 0.04331086, -0.23139133, -0.09555604, 0.06358018, + -0.59850585, -0.37260893, -0.9738547 , -0.01462804, 0.12405673, + 0.07412421, -0.01379338, 0.05929387, 0.7805736 , -0.36394337, + 0.47818112, -0.43866792, 0.63309664, 0.1147754 , 0.45001298, + 0.3213029 , 0.44093522, 0.51016414, -0.18746474, -0.13791506, + 0.04778931, -0.49275073, 0.04839996, -0.38444966, -0.4882343 , + -0.6615728 , 0.09630081, -0.29930323, -0.34863642, 0.77798736, + -0.55292934, -0.56388927, 0.09340773, -0.4621951 , 0.26494777, + 0.31161246, 0.44366807, 0.6213311 , 0.01337955, 0.14917652, + -0.4876238 , 0.06336308, -0.12776731, -0.07288839, -0.18272795, + 0.23383974, -0.059283 , -0.08504926, -0.12375498, 0.3620917 , + 0.04436552, 0.15867496, 0.12362888, -0.09505215, 0.18897517, + 0.08316641, -0.18332016, -0.26461446, 0.17972536, 0.23397446, + -0.04752706, -0.04319599, 0.12071403, 0.25849587, 0.06566656, + 0.19910227, 0.22742352, -0.22645345, -0.14635295, -0.2507506 , + -0.18631294, -0.05247986, 0.12347726], dtype=float32), + kernel: Array([[-0.06023952, 0.297179 , -0.02699618, ..., -0.1145321 , + -0.08048212, -0.10587901], + [ 0.6934901 , 0.47458088, 0.61580837, ..., -0.13140583, + -0.14733373, 0.26471645], + [ 0.09990517, -0.17357504, 0.05461919, ..., -0.11821933, + -0.24399522, -0.15253429], + ..., + [-0.2205276 , -0.23215702, -0.24366312, ..., -0.03175567, + 0.22550093, 0.27891138], + [-0.24111255, -0.27693352, -0.21988523, ..., 0.30499363, + 0.37258562, -0.13957877], + [-0.103215 , -0.14666593, -0.04492585, ..., -0.6473914 , + -0.23865867, 0.19981082]], dtype=float32), + }, + query: { + bias: Array([-0.11081476, -0.09144565, -0.06050301, -0.03483055, 0.03535065, + -0.1849559 , -0.04721399, -0.07624152, -0.10357637, 0.22812082, + -0.3447222 , 0.0331145 , 0.05764522, -0.13265188, 0.1981749 , + -0.03139308, -0.07606025, 0.13599226, 0.02934942, 0.12725802, + 0.07556148, 0.07594807, -0.04966472, -0.15265177, -0.35979733, + 0.01576409, -0.07025045, -0.19974303, -0.22990641, 0.0952646 , + -0.33106107, 0.04587565, -0.41658264, 0.33157024, -0.51854265, + 0.34642565, -0.38225 , 0.21841964, 0.4542027 , -0.3708585 , + -0.17867757, 0.38574484, -0.40515402, -0.18270488, 0.15764216, + -0.09051255, -0.47678202, -0.5485276 , 0.17880292, -0.42515096, + 0.11444188, 0.27703342, -0.23375264, -0.17498691, 0.13194664, + 0.06895838, 0.32982063, 0.17495483, -0.16514482, 0.16266893, + 0.24224894, 0.49473688, -0.7105683 , -0.22717172, -0.4628581 , + -0.02517562, 0.13951837, 0.04581514, -0.02814059, 0.06052845, + 0.89534795, -0.01086807, -0.2937633 , 0.4263962 , 0.01343843, + 0.9783011 , 0.23017581, 0.7163039 , 0.41632223, 0.17262603, + -0.15068051, 0.850449 , -0.16111681, -0.02732765, -0.18101065, + 0.47148392, 0.18522678, -0.06222672, 0.15925601, -0.07961265, + 0.71117395, 0.02433823, -0.35966122, -0.09352748, -0.06117439, + -0.13183966, -0.08041007, -0.27098748, -0.71445733, -0.5658294 , + 0.51616377, -0.56782836, -0.68260497, -0.5732899 , 0.5585706 , + -0.14336808, 0.59159863, -0.21151459, 0.42952493, 0.46055183, + 0.46279824, -0.4195606 , -0.5513547 , 0.34171253, 0.69069153, + -0.25621027, -0.46899697, 0.588077 , 0.6255269 , 0.5819053 , + 0.54691887, 0.59542644, -0.5749214 , -0.57993734, 0.42467687, + -0.70540833, -0.34917867, 0.3106283 ], dtype=float32), + kernel: Array([[-0.11731137, -0.1564569 , -0.11458346, ..., 0.02826559, + -0.03335454, 0.14976554], + [-0.04609706, -0.04719561, -0.02786452, ..., -0.30554634, + -0.20783764, -0.11992735], + [ 0.0489803 , 0.12078558, 0.0699265 , ..., 0.26866522, + -0.40512353, -0.5030211 ], + ..., + [-0.00423968, -0.00703923, -0.03773196, ..., -0.10513086, + -0.145968 , 0.07945077], + [ 0.07510082, -0.13710845, 0.03422854, ..., 0.15533802, + 0.8349194 , -0.10965192], + [-0.07718561, 0.09582806, -0.06439573, ..., -0.12381844, + -0.9082412 , 0.04665521]], dtype=float32), + }, + value: { + bias: Array([-0.11375186, 0.01737056, -0.11132707, -0.01888963, 0.02005741, + -0.00092973, -0.11860004, 0.14434907, -0.06270761, 0.14292052, + 0.14930397, 0.06040736, 0.1552559 , -0.06054825, 0.13105005, + 0.05230521, 0.03080181, -0.14906318, 0.04209505, -0.11033 , + -0.02365249, -0.14162254, -0.00175851, 0.03008844, 0.04629046, + 0.10597341, 0.02212737, 0.13431458, -0.03534039, 0.00220095, + -0.05828169, 0.01887055, 0.07417873, -0.18334955, 0.05120145, + 0.03194632, 0.09620124, 0.07511445, 0.14970838, -0.3099061 , + -0.00785524, 0.21368372, -0.24593364, -0.20003676, 0.29439747, + 0.02457003, -0.0362221 , -0.04669455, -0.11111037, -0.09324432, + -0.007548 , -0.05471316, -0.1160496 , -0.08286496, 0.10508952, + -0.29490575, 0.00155351, 0.13688038, 0.055782 , -0.241635 , + 0.19978897, 0.12800941, -0.08591252, -0.17519936, 0.0762416 , + 0.03798862, -0.04639675, -0.1959964 , -0.25277418, 0.169555 , + -0.06163629, -0.11278319, -0.01159456, -0.06765992, 0.0005254 , + -0.06396206, -0.00113279, 0.1249745 , -0.02429713, -0.01882967, + -0.04579956, -0.02704 , 0.00707008, 0.02282375, -0.10874218, + -0.01254413, 0.05641572, -0.06401654, 0.17181516, 0.0763936 , + -0.07932342, -0.12056571, 0.0820564 , -0.01728764, 0.04256084, + -0.05678854, 0.11785327, 0.2377933 , -0.19259618, -0.08974701, + -0.19949402, -0.09832341, 0.12747605, 0.01983653, 0.00941677, + 0.11338363, -0.20597233, -0.04107258, 0.02300076, -0.04743035, + 0.11633945, 0.12718917, 0.23441505, 0.01317911, -0.19377421, + -0.06671913, 0.04473834, 0.00467375, -0.15131958, -0.07110944, + -0.16257668, -0.08920506, 0.03026637, -0.0935519 , 0.06253995, + -0.19525133, -0.21808124, -0.19407861], dtype=float32), + kernel: Array([[-0.01158661, 0.04303771, -0.05590956, ..., -0.05333043, + -0.00572822, -0.01724763], + [ 0.2719801 , 0.14965016, -0.02898863, ..., -0.06263707, + -0.35550225, -0.28345698], + [-0.05632893, -0.26717782, -0.07165924, ..., -0.43430653, + 0.0922191 , 0.10500429], + ..., + [-0.04451123, -0.00324844, 0.15866533, ..., -0.15973136, + -0.06277192, -0.10269681], + [-0.28674033, 0.40732047, -0.23387024, ..., 0.5245572 , + -0.20356956, -0.19302332], + [-0.02055671, -0.35077375, 0.13905095, ..., 0.08351879, + 0.4497927 , 0.44072792]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-2.29397908e-01, 3.42801064e-01, 1.12294301e-01, -1.40454136e-02, + 3.03802863e-02, -1.64602742e-01, 6.72679991e-02, 9.77464989e-02, + -2.96701267e-02, 7.89122283e-02, 5.59587255e-02, 1.48079200e-02, + -1.85802028e-01, -3.27114910e-02, -2.40307394e-02, 6.82050884e-02, + 5.04146097e-04, -1.56787038e-01, 6.75788894e-02, -1.31789781e-02, + -4.55502272e-02, 1.27612054e-01, -7.01643378e-02, -2.48179883e-02, + -1.31983072e-01, 1.38503700e-01, 1.17344763e-02, 7.56531805e-02, + 6.57847598e-02, 2.69352719e-02, 2.21325293e-01, -5.05222380e-02, + 3.61935459e-02, -1.04897447e-01, 5.24883904e-02, -8.26861486e-02, + 1.44010643e-02, 8.25399533e-02, -6.71743648e-03, -1.94458306e-01, + -2.49852352e-02, -8.23389366e-02, -6.81263685e-01, 8.11951086e-02, + -8.49441960e-02, -1.40838427e-02, -1.44380983e-02, 6.43714517e-02, + -1.57576092e-02, -3.63134891e-02, -1.25548737e-02, 1.41190467e-02, + -1.39931515e-01, 1.46125898e-01, -4.10604011e-03, -3.28917056e-02, + -1.47544041e-01, -7.30510205e-02, 1.31265923e-01, 3.08194682e-02, + 5.00784516e-02, -2.18488395e-01, 2.87581724e-03, -3.31245400e-02, + 2.70526409e-02, -1.38555646e-01, -1.16920797e-02, -7.85453171e-02, + 1.44379027e-02, -8.10271129e-02, -7.89678767e-02, 6.87262276e-03, + 2.95094848e-02, -2.96989549e-02, 8.52790475e-01, 1.68656245e-01, + -4.95392680e-02, 1.01274252e-03, -1.17389508e-01, 1.42783284e-01, + -2.95234203e-01, -3.66869681e-02, 2.10161004e-02, 1.61923394e-01, + 1.00733116e-01, 9.56858695e-02, 1.62279665e-01, 6.97396919e-02, + -6.76168874e-02, -1.21831730e-01, 2.96244565e-02, -8.43596309e-02, + -4.85159941e-02, -1.54231772e-01, -8.46877992e-02, -9.01621133e-02, + -7.94285350e-03, 9.12197307e-03, -5.20062111e-02, 2.70500332e-02, + 1.75196305e-01, -7.51989707e-02, 1.59212708e-01, -8.15638751e-02, + 1.15850195e-01, -1.62196823e-03, 6.47357618e-03, -1.16396412e-01, + -1.10036433e-01, -1.06435921e-02, -1.69102713e-01, 2.04750933e-02, + 1.00554541e-01, -3.42080109e-02, 1.30484760e-01, 1.22546069e-02, + -1.68031573e-01, 3.20073776e-02, -6.57493770e-02, -7.96901435e-03, + 1.29959419e-01, 8.22199881e-02, 4.45892625e-02, -4.60687764e-02, + 6.53466284e-02, 2.84257364e-02, -7.09084561e-03, 8.53643045e-02], dtype=float32), + kernel: Array([[-4.33085859e-01, -5.61205745e-01, 2.93745458e-01, ..., + -1.06319696e-01, -3.59772116e-01, -5.50748073e-02], + [ 1.29990995e-01, -1.94321834e-02, 3.05669427e-01, ..., + 1.81368724e-01, 9.14699510e-02, -6.21524341e-02], + [-2.22445935e-01, 1.21082552e-01, -5.40154229e-04, ..., + -1.46840677e-01, -1.43376559e-01, 9.97718945e-02], + ..., + [ 6.94694817e-01, 8.00504535e-03, 7.26794839e-01, ..., + -2.44273514e-01, -3.41350943e-01, -1.44506514e-01], + [ 6.64600372e-01, -1.74710259e-01, 2.86397934e-01, ..., + -4.65431027e-02, -2.18760711e-03, -2.81117141e-01], + [ 5.85957587e-01, -1.68455601e-01, 3.16843003e-01, ..., + -4.89029028e-02, -3.35241156e-03, -2.94326484e-01]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([ 6.88922033e-02, 5.80978803e-02, 3.85873437e-01, -3.11470646e-02, + 1.69555983e-03, -2.32034996e-01, -1.58923380e-02, -3.62700596e-02, + 5.75650949e-04, -3.82705927e-01, 2.31035739e-01, 6.68624789e-02, + -1.72114998e-01, -5.66142239e-02, 2.09674109e-02, -2.80936039e-03, + -2.27225497e-02, -3.59404653e-01, 1.31325588e-01, -3.94533155e-03, + -1.76445603e-01, 4.06457812e-01, -3.71316727e-03, -2.46738438e-02, + -1.81620270e-01, 4.64766249e-02, 7.05432147e-02, 7.28050619e-02, + -5.19782789e-02, -4.45098169e-02, 3.48823667e-01, 6.42984435e-02, + -3.04025542e-02, -9.68929306e-02, -5.14351986e-02, -3.97855928e-03, + -3.20420787e-02, 2.67196506e-01, -7.22298101e-02, -2.07284018e-01, + 4.90039028e-02, -2.19781891e-01, -1.35466352e-01, 8.54400694e-02, + 1.19789563e-01, -1.32402414e-02, -2.06071258e-01, 3.45563710e-01, + 1.00530580e-01, 8.49714726e-02, 3.91122419e-03, 2.15060711e-02, + 8.78436714e-02, 2.07931370e-01, 3.97193991e-02, -1.13610223e-01, + 3.73787954e-02, 1.23064490e-02, -9.43080038e-02, -1.86623499e-01, + -3.85700941e-01, -1.63304899e-02, 1.10437632e-01, -1.03100136e-01, + -6.19614311e-03, -3.90111618e-02, 1.53294563e-01, -1.00014314e-01, + 8.98484327e-03, -1.33218288e-01, 1.19500153e-01, -1.21143341e-01, + 2.10582778e-01, 1.27065793e-01, -1.65820807e-01, 2.54214227e-01, + 5.12329042e-02, -9.36500803e-02, 5.69578595e-02, 3.65398824e-01, + -1.66889839e-02, -8.60542282e-02, -2.17619090e-04, 7.38569945e-02, + 1.74466029e-01, 7.53093883e-02, -1.50982007e-01, -1.61455601e-01, + -4.24303383e-01, -1.41928092e-01, 1.67686343e-02, -2.84632683e-01, + 1.45200789e-01, 5.28667821e-03, -2.16334872e-02, 4.74276505e-02, + 1.93778221e-02, -1.02245426e-02, 8.27698484e-02, -4.25329745e-01, + -9.56311524e-02, 2.45452020e-02, 5.58077872e-01, 8.93760026e-02, + -3.42835188e-02, -7.27823749e-03, -4.62576747e-02, -1.11996412e-01, + -2.86399245e-01, 3.93988378e-03, -5.40265851e-02, 1.06320893e-02, + -1.11137971e-01, 9.71371382e-02, 7.98522159e-02, -7.13857710e-02, + -4.83553410e-01, 1.70313239e-01, -1.70816466e-01, 8.59201849e-02, + -3.64008136e-02, -2.24960689e-02, -3.12829800e-02, 1.58451021e-01, + 9.86519852e-04, -3.24616969e-01, -2.68537730e-01, -1.29793763e-01], dtype=float32), + scale: Array([ 0.16587867, 0.62729955, 0.78539467, 0.33312872, 0.02448054, + 0.52365166, 0.09352965, 0.57490396, -0.00181935, 0.62166315, + 0.5622861 , 0.632528 , 0.32434413, 0.33608988, 0.09364542, + 0.00295282, 0.20884101, 0.4969771 , 0.49006283, 0.3390643 , + 0.76553243, 0.75735354, 0.64736444, 0.20501752, 0.5099314 , + 0.39772737, 0.4760952 , 0.2673762 , 0.5248422 , 0.34622046, + 0.67513955, 0.12790139, 0.78136003, 0.46376252, 0.17142247, + 0.02035132, 0.35290378, 0.79245657, 0.43180364, 0.58383363, + 0.28102416, 0.6248626 , 0.09244277, 0.38200957, 0.36086422, + 0.44281945, 0.62160456, 0.49688524, 0.47265998, 0.4368852 , + 0.00439632, 0.24748251, 0.2077039 , 0.71143705, 0.55966383, + 0.76935214, 0.2611507 , 0.3572642 , 0.4780699 , 0.5534745 , + 0.5650403 , 0.49476576, 0.6574258 , 0.38081974, 0.15224549, + 0.2871514 , 0.5269491 , 0.5443608 , 0.47709846, 0.6731636 , + 0.48256594, 0.57757723, 0.52946705, 0.27599993, 0.11646109, + 0.7459179 , 0.6025415 , 0.4426929 , 0.150529 , 0.6974664 , + 0.07860193, 0.63189274, 0.6666647 , 0.8469711 , 0.4019909 , + 0.45879367, 0.22790541, 0.5310602 , 0.46144342, 0.46353984, + 0.3125676 , 0.6395497 , 0.5902147 , -0.00431166, 0.51785195, + 0.44832194, 0.2982607 , -0.02292351, 0.48743743, 0.42023692, + 0.5205952 , 0.08753084, 0.76318944, 0.42676944, 0.6018406 , + 0.7163417 , 0.55152804, 0.8593687 , 0.4082306 , 0.4901796 , + 0.64518034, 0.5525667 , 0.34930795, 0.46910506, 0.6457558 , + 0.54612505, 0.61357117, 0.49350467, 0.41991618, 0.55927193, + 0.4409146 , 0.15714881, 0.31889114, 0.26255667, 0.45123503, + 0.56282735, 0.50075245, 0.34856102], dtype=float32), + }, + layernorm_before: { + bias: Array([ 2.82420870e-02, -2.29402691e-01, -2.96053607e-02, -1.30823404e-02, + -2.65389867e-03, 2.99529568e-03, -3.38511716e-04, -4.37224060e-02, + 2.60172427e-01, -2.05475893e-02, 1.60881188e-02, -1.69427078e-02, + -5.21531154e-04, 7.41502224e-03, -1.61934067e-02, 3.54598239e-02, + -1.85820572e-02, 2.71367803e-02, -4.88283485e-03, -6.32894114e-02, + 3.68127339e-02, -6.04641587e-02, 1.19915558e-03, 1.77571445e-03, + -7.46919066e-02, -1.53914511e-01, -3.68062826e-03, 1.60319414e-02, + 6.28713518e-02, 3.06911282e-02, -2.31577735e-02, -6.57155272e-03, + 6.15396686e-02, 2.63886862e-02, -4.85325791e-02, 7.88616296e-03, + 8.01346614e-04, -1.55313745e-01, -2.11913697e-03, -6.54404704e-03, + 5.55190295e-02, -7.31622195e-03, 1.96012199e-01, 5.94329368e-03, + -2.78697144e-02, -4.93782759e-02, 3.01586390e-02, 6.05354905e-02, + -1.21998333e-03, 3.04064993e-02, 1.04687298e-02, 1.09375557e-02, + -2.71287002e-03, -2.32986733e-01, 1.44290384e-02, -2.30109710e-02, + 4.33547981e-02, 5.82204685e-02, -2.50215475e-02, 1.38790637e-01, + -6.68904185e-02, 1.24691881e-01, -4.73514060e-03, -3.82121443e-03, + 3.83599102e-02, 8.80813003e-02, -1.65379103e-02, -7.19220552e-04, + 3.60673331e-02, 1.00804691e-03, -1.36750117e-02, -2.13624742e-02, + 3.65178101e-02, 2.25900616e-02, -4.44582880e-01, -2.65965145e-02, + -5.42465337e-02, -7.62177110e-02, 6.57133060e-04, -1.34424558e-02, + 3.63113172e-02, 7.55314380e-02, 8.09766643e-04, -2.51245238e-02, + -2.74020876e-03, -4.82323244e-02, -7.02568442e-02, 1.68840366e-03, + 1.02340458e-02, 1.46568000e-01, -2.30909362e-02, 4.56218012e-02, + -6.14786409e-02, -2.31754649e-02, 1.12327132e-02, 8.98616686e-02, + 3.33565250e-02, 3.43459705e-03, -1.15024485e-02, -9.24066082e-02, + -3.95509861e-02, -2.51406804e-03, -9.44837704e-02, -6.91622263e-04, + -3.78696583e-02, -7.13908523e-02, -6.70393333e-02, 2.14871876e-02, + -3.32830548e-02, -9.49056670e-02, 3.31703201e-02, 5.51439868e-03, + -1.63990129e-02, 1.68498412e-01, -1.52823254e-01, 1.28563598e-01, + 2.69770622e-02, -2.52538477e-03, -6.20844103e-02, -2.31217127e-03, + -1.12146828e-02, 2.15321942e-03, -1.07783070e-02, 7.04224780e-03, + -5.70042757e-04, -1.27387140e-02, -1.50637852e-03, -1.26768842e-01], dtype=float32), + scale: Array([ 0.02225583, 0.3582427 , 0.18810613, 0.2897796 , 0.01008445, + 0.20093079, 0.00262845, 0.18910973, 0.20207141, 0.05451531, + 0.13650979, 0.22285123, 0.04530312, 0.00609608, 0.10985828, + 0.09655713, 0.13278736, 0.08787937, 0.02501552, 0.16963154, + 0.22782324, 0.20419084, 0.23176458, 0.04793768, 0.12091924, + 0.15503103, 0.12642743, 0.07698647, 0.37561363, 0.1734934 , + 0.22910334, 0.00436596, 0.25682634, 0.24328624, 0.2714059 , + 0.06567708, 0.02274263, 0.567844 , -0.00143492, 0.16284508, + 0.18691777, 0.3501844 , 0.11066452, 0.01387185, 0.17211044, + 0.31377098, 0.1559427 , 0.31400132, 0.03055096, 0.1480449 , + 0.06916483, 0.05761943, -0.00319004, 0.38499779, 0.0325634 , + 0.3114719 , 0.07862399, 0.19379751, 0.12478653, 0.41179135, + 0.26195255, 0.3847175 , 0.16452418, 0.00272222, 0.15891619, + 0.11142635, 0.14923933, 0.10400768, 0.21697171, 0.12220434, + 0.24840005, 0.16391885, 0.09268359, 0.12174353, 0.12482438, + 0.2716236 , 0.22829536, 0.2231431 , -0.0092565 , 0.03284183, + 0.06326266, 0.15956563, 0.06667496, 0.22676413, 0.01954163, + 0.2572471 , 0.14380816, 0.3025213 , 0.13769881, 0.38677344, + 0.05833767, 0.11660744, 0.1784911 , 0.31127843, 0.0378647 , + 0.11842007, 0.13203348, 0.08775622, 0.22929348, 0.1685133 , + 0.17730162, 0.14410907, 0.3452061 , -0.01377543, 0.44221473, + 0.43522978, 0.29845122, 0.13120209, 0.45809656, 0.19538262, + 0.46633697, 0.1401029 , 0.05973453, 0.4215374 , 0.16866335, + 0.27390364, 0.03285838, 0.01355824, 0.18128489, 0.2254643 , + 0.03559769, 0.00747606, 0.18407959, 0.02931737, -0.00183865, + 0.0155194 , 0.2910443 , 0.2715314 ], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-2.87663639e-01, -2.04487324e-01, -3.44864547e-01, -2.00529769e-01, + -5.94829321e-01, -1.03927039e-01, -2.77930766e-01, -4.94999826e-01, + -2.08718747e-01, -3.45035255e-01, -2.67087758e-01, -4.77189213e-01, + -1.66359320e-01, -4.48330075e-01, -4.78699297e-01, -2.21549183e-01, + -2.41500542e-01, -3.03187490e-01, -4.03688133e-01, -3.62672597e-01, + 3.11330687e-02, -4.75372255e-01, -1.33399740e-01, -1.58057332e-01, + -1.38183832e-01, -3.53402674e-01, -4.83101279e-01, -3.49856049e-01, + -4.62018728e-01, -2.20034316e-01, -1.45977482e-01, -1.44924179e-01, + -3.43084693e-01, -2.04671875e-01, -2.16533989e-01, -5.23977339e-01, + -2.12659031e-01, -2.74191111e-01, -2.65450865e-01, -4.38724220e-01, + -4.05285835e-01, -4.25657749e-01, -2.83261627e-01, -1.07905425e-01, + -3.10509741e-01, -3.69468153e-01, -5.25859118e-01, -3.74539018e-01, + -4.44122553e-01, -2.76227951e-01, -3.62340897e-01, -3.20163608e-01, + -5.22766829e-01, -3.68136406e-01, -6.11795008e-01, -1.88745782e-01, + -4.54058141e-01, -1.91686556e-01, -4.25439298e-01, -5.22909224e-01, + -3.05993110e-01, -3.57376486e-01, -3.10534716e-01, -2.63750911e-01, + -5.45675814e-01, -4.23284829e-01, -1.21132560e-01, -3.99448633e-01, + -1.06209874e-01, -3.28158796e-01, -5.71545482e-01, -1.91518441e-01, + -5.93747199e-01, -3.57770473e-01, -2.59463459e-01, -3.54470700e-01, + -5.28468788e-01, -3.18399936e-01, -4.58652407e-01, -2.15082556e-01, + -4.01729345e-01, -5.17883539e-01, -3.26544613e-01, -2.71224737e-01, + -3.98273647e-01, -3.58481169e-01, -4.33161110e-01, -4.17540729e-01, + -4.41764146e-01, -4.13900316e-01, -4.20035869e-01, -4.61230427e-01, + -5.54932415e-01, -2.82214522e-01, -4.63347554e-01, -1.69810832e-01, + -4.96465772e-01, -3.57517421e-01, -3.08119923e-01, -5.05688906e-01, + -4.42314088e-01, -2.96192914e-01, -3.34470212e-01, -1.58352673e-01, + -1.96208924e-01, -1.71411768e-01, -2.78193712e-01, -4.14353907e-01, + -5.30355275e-01, -6.04073517e-02, -3.42649519e-01, -3.25455546e-01, + -2.80562103e-01, -3.16360086e-01, -1.88339218e-01, -1.45719707e-01, + -2.42292494e-01, -4.38852787e-01, -3.38206112e-01, -3.96864712e-01, + -3.94322276e-01, -5.04216313e-01, -4.47647184e-01, -3.14450145e-01, + -4.03099447e-01, -3.31715465e-01, -5.11624277e-01, -2.97003418e-01, + -1.71126753e-01, -5.23082614e-01, -3.97670537e-01, -3.33513230e-01, + -2.73618609e-01, -2.16470793e-01, -3.16209495e-01, -3.54464471e-01, + -4.11175519e-01, -3.47123802e-01, -4.30214375e-01, -2.96357393e-01, + -3.50862324e-01, -4.91615891e-01, -2.99721777e-01, -3.46929997e-01, + -6.17211461e-01, -3.17567289e-01, -2.59017050e-01, -5.43685555e-01, + -3.07671517e-01, -3.49160552e-01, -5.83933473e-01, -2.99324632e-01, + -3.56959522e-01, -4.75278974e-01, -4.69556153e-01, -4.00068551e-01, + -3.57878387e-01, -2.61400610e-01, 5.89218475e-02, -4.75838035e-01, + -3.18645447e-01, -2.91967630e-01, -3.57156008e-01, -5.22750139e-01, + -1.99007168e-01, -4.69590366e-01, -3.96272272e-01, -2.64699906e-01, + -3.60923529e-01, -1.91638634e-01, -3.21910053e-01, -2.48550743e-01, + -3.98402393e-01, -4.65990812e-01, -3.43011886e-01, -3.76989871e-01, + -3.99559021e-01, -3.55907410e-01, -4.88500834e-01, -3.14257860e-01, + -5.45107484e-01, -1.72966227e-01, -3.69181097e-01, -2.42310494e-01, + -3.32116067e-01, -3.86646032e-01, -2.56473958e-01, -3.48980218e-01, + -3.98971230e-01, -2.64990151e-01, -4.63962048e-01, -1.67646125e-01, + -3.03776145e-01, -2.27759063e-01, -4.02701110e-01, -3.71156693e-01, + -4.49759185e-01, -4.83259469e-01, -4.44558293e-01, -3.49739552e-01, + -4.34326917e-01, -2.64975369e-01, -2.89425910e-01, -4.49308783e-01, + -2.49016702e-01, -4.31563139e-01, -4.08892900e-01, -3.69399726e-01, + -6.46372512e-02, -3.25122327e-01, -1.55182466e-01, -1.11164600e-01, + -4.79935288e-01, -5.89126289e-01, -5.05477190e-01, -3.91809165e-01, + -4.87700224e-01, -5.42157114e-01, -5.01625121e-01, -3.88096064e-01, + -4.09615874e-01, -3.30858119e-03, -3.16691220e-01, -4.25819814e-01, + -2.07181290e-01, -2.52645940e-01, -4.78967577e-01, -4.66646552e-01, + -5.08263230e-01, -2.38753840e-01, -3.27672154e-01, -5.24233878e-01, + -3.54793280e-01, -3.08880091e-01, -2.62713552e-01, -3.00624251e-01, + -4.21428025e-01, -4.42182183e-01, -5.11559665e-01, -2.42620170e-01, + -3.54224473e-01, -3.06683362e-01, -4.52630949e-04, -4.44990546e-01, + -3.27576101e-01, -2.59946615e-01, -5.11295557e-01, -6.64835870e-01, + -1.85168549e-01, -5.48847735e-01, -2.54403502e-01, -4.44106311e-01, + -2.09086508e-01, -5.46322465e-01, -4.56374437e-01, -2.26774529e-01, + -3.03409159e-01, -2.95624912e-01, -4.80216295e-01, -1.95244759e-01, + -2.66378999e-01, -3.54911506e-01, -3.88011873e-01, -2.92912245e-01, + -4.73809302e-01, -2.48538107e-01, -2.53374189e-01, -5.34175217e-01, + 1.01788407e-02, -4.37775195e-01, -3.74520659e-01, -4.16397095e-01, + -2.86437750e-01, -1.17193200e-01, -4.65466827e-01, -2.45054185e-01, + -1.55902386e-01, -2.60219306e-01, -3.81529987e-01, -4.36873585e-01, + -3.58277500e-01, -4.33071882e-01, -4.30247694e-01, -3.75352204e-01, + -3.81595612e-01, -4.86366630e-01, -4.01012272e-01, -2.53684759e-01, + -3.21343958e-01, -5.21693230e-01, -2.32155509e-02, -3.01609904e-01, + -2.36861959e-01, -3.18232119e-01, -1.21864028e-01, -3.91450346e-01, + -5.88354170e-01, -3.05314749e-01, -3.78017843e-01, -9.56703201e-02, + -3.85766625e-01, -2.07299024e-01, -8.47784728e-02, -3.33295763e-01, + -4.15781707e-01, -4.48703825e-01, -1.40794486e-01, -1.18951991e-01, + -2.74386436e-01, -4.96934414e-01, -3.64418119e-01, -3.70397896e-01, + -1.35903224e-01, -1.35731950e-01, -3.99609268e-01, -2.93618023e-01, + -3.04051191e-01, -5.70618927e-01, -5.02060175e-01, -2.81283766e-01, + -6.31259501e-01, -3.06242526e-01, -2.12949291e-01, -4.86519247e-01, + -1.59262896e-01, -4.03478205e-01, -1.81495771e-01, -4.62975919e-01, + -3.86971891e-01, -5.12782514e-01, -2.76577592e-01, -2.94954062e-01, + -2.17922330e-01, -4.69827503e-01, 1.06356489e-02, -3.15419227e-01, + -3.53538573e-01, -5.48389792e-01, 5.30533865e-02, -3.23666394e-01, + -3.68639082e-01, -3.19764733e-01, -4.70012695e-01, -9.86519307e-02, + -3.48372728e-01, -2.50590980e-01, -3.00890654e-01, -4.49694991e-01, + -3.42396975e-01, -2.40296051e-01, -1.83449000e-01, -4.94159639e-01, + -3.83832961e-01, -2.35955358e-01, -4.39356416e-01, -2.83128679e-01, + -5.06890595e-01, -4.84923810e-01, -3.48398417e-01, -8.36336315e-02, + -3.36799473e-01, -4.19064090e-02, -5.25771677e-01, -4.51774269e-01, + -4.36057091e-01, -4.80761558e-01, -3.56559277e-01, -4.32838291e-01, + -4.25217807e-01, -2.08370864e-01, -6.56363890e-02, -5.52178204e-01, + -4.09472764e-01, -1.91283375e-01, -1.69247791e-01, -3.70922804e-01, + -4.36929941e-01, -1.94023415e-01, -2.98483998e-01, -1.78429320e-01, + -2.56660700e-01, -3.49405706e-01, -2.60902107e-01, -1.27692908e-01, + -2.62808502e-01, -3.31191778e-01, -5.31970501e-01, -2.73841202e-01, + -3.46037507e-01, -2.07595572e-01, -3.45578700e-01, -3.25063765e-01, + -3.59042734e-01, -5.38096130e-01, -3.52915227e-01, -2.89237618e-01, + -4.53390926e-01, -3.57132286e-01, -2.32824370e-01, -1.73414394e-01, + -3.02039832e-01, -4.06891137e-01, -3.11633587e-01, -3.86790335e-01, + -4.44779336e-01, -4.03608739e-01, -3.27035785e-01, -5.77829838e-01, + -3.05345178e-01, -5.01044929e-01, -4.90038365e-01, -3.42383653e-01, + -2.16015473e-01, -2.34810472e-01, -3.94850314e-01, -3.71235818e-01, + 1.26175291e-03, -4.32831019e-01, -1.67813487e-02, -4.22147512e-01, + -3.73436451e-01, -4.40798849e-01, -4.10787642e-01, -4.72505748e-01, + -2.18411952e-01, -3.72730881e-01, -3.14533114e-01, -2.04394564e-01, + -2.78175622e-01, -2.42158204e-01, -3.91372234e-01, -5.00931203e-01, + -2.58081257e-01, -4.08814311e-01, -4.02594358e-02, -3.35627794e-02, + -2.13214949e-01, -3.36978287e-01, -2.62033820e-01, -6.90276861e-01, + -3.95373791e-01, -4.39152598e-01, -2.85017580e-01, -4.47889686e-01, + -4.36591983e-01, -3.10838968e-01, -4.41035658e-01, -2.79656291e-01, + -1.40883625e-01, -3.65119666e-01, -1.61859736e-01, -1.76689908e-01, + -3.66094023e-01, -4.24499989e-01, -4.19989884e-01, -2.08415791e-01, + -2.32176647e-01, -3.55920881e-01, -2.76795477e-01, -4.64629948e-01, + -2.86160856e-01, -5.59463799e-01, -3.76166940e-01, -3.82688463e-01, + -4.46783721e-01, -3.94609421e-01, -3.84682864e-01, -2.80228592e-02, + -3.13305557e-01, -3.56978267e-01, -4.11862016e-01, -6.72479212e-01, + -2.63180643e-01, -2.53385037e-01, -4.69251305e-01, -4.94098127e-01, + -3.67969900e-01, -5.01329899e-01, -4.63324428e-01, -4.17908072e-01, + -2.42769625e-03, -3.66800427e-01, -4.42783326e-01, -5.09920537e-01, + -3.79778743e-01, -2.11185992e-01, -4.76333320e-01, -2.97944933e-01, + -5.35430849e-01, -5.54302871e-01, -1.78685069e-01, -5.05887330e-01, + -3.60490650e-01, -2.97754109e-02, -4.16262656e-01, -2.02613860e-01, + -4.54910725e-01, -5.27239859e-01, -3.19837511e-01, -3.31910402e-01, + -5.04021525e-01, -5.96606016e-01, -2.31208146e-01, -5.22311389e-01, + -3.61870646e-01, -2.52372891e-01, -5.43789089e-01, -3.68237197e-01, + -2.00233698e-01, -3.31239879e-01, -3.89182925e-01, -1.77743852e-01], dtype=float32), + kernel: Array([[-0.03629679, 0.12771748, 0.30917847, ..., -0.02645334, + 0.2853967 , 0.25290728], + [-0.05366539, -0.08199085, -0.77184826, ..., -0.28385365, + -0.08720582, -0.2480316 ], + [-0.30405623, -0.14440256, -0.16176778, ..., -0.6061087 , + -0.53592414, -0.557683 ], + ..., + [-0.5002675 , 0.6923438 , -0.39955443, ..., 0.2793648 , + -0.0890847 , 0.13734408], + [ 0.42013338, -0.1837182 , 0.20519936, ..., 0.59244996, + 0.15519385, -0.04555327], + [ 0.2887559 , 0.03648978, 0.19605577, ..., -0.3638592 , + -0.05200996, 0.14020972]], dtype=float32), + }, + output: { + bias: Array([-2.03454956e-01, 2.89009154e-01, -5.26833301e-03, 1.99230332e-02, + 7.02989474e-03, -8.10585693e-02, 1.01039574e-01, 1.28230199e-01, + 1.65420137e-02, 1.76779047e-01, 2.31052637e-02, -2.84619331e-02, + -1.72500789e-01, 8.57431665e-02, -2.58707665e-02, 8.68742466e-02, + 3.48051847e-03, -8.32902640e-02, 8.47522356e-03, -6.26579765e-03, + 6.64773211e-02, -5.39215617e-02, -1.33164413e-02, -2.61053070e-02, + -1.25059590e-01, 9.69652552e-03, -6.74541369e-02, 5.93836829e-02, + 6.43064156e-02, 6.59576431e-02, 1.18919395e-01, -1.21684916e-01, + 1.34139672e-01, -7.79727027e-02, 7.02831373e-02, -7.30708465e-02, + 2.56273188e-02, 8.22603703e-03, 5.56182116e-03, -1.27773434e-01, + -1.63553078e-02, -2.40529538e-03, -7.88959324e-01, 2.98491977e-02, + -1.15960285e-01, -1.11549925e-02, 1.14158385e-01, 4.73596305e-02, + -4.17886041e-02, -5.26564419e-02, -9.78346448e-03, 4.26121093e-02, + -1.87094569e-01, 8.54555890e-02, -4.69494089e-02, -3.20936963e-02, + -1.27747595e-01, -3.98493744e-02, 1.56673312e-01, 5.93540259e-02, + 1.29704997e-01, -2.26318836e-01, -4.20836620e-02, 1.89136285e-02, + 2.25093830e-02, -1.10592782e-01, -8.14836994e-02, -1.74993854e-02, + 6.80994093e-02, -1.73572097e-02, -1.06341034e-01, 2.98510771e-03, + -6.83057960e-03, -4.16511074e-02, 8.99981439e-01, 1.04607493e-02, + -1.10422350e-01, 4.36907373e-02, -1.15947202e-01, -4.37050452e-03, + -4.16922748e-01, 3.20061632e-02, 4.93779741e-02, 2.31584487e-03, + -1.06097618e-02, 8.79805684e-02, 1.83927685e-01, 1.49958774e-01, + 1.10785790e-01, 1.96547806e-02, -6.48459792e-03, 2.37252992e-02, + -1.06167868e-01, -1.58290878e-01, -3.44699137e-02, -5.64344190e-02, + 2.31592474e-03, 1.24252474e-04, -6.33904487e-02, 1.05357006e-01, + 2.25357682e-01, -6.24567755e-02, 3.09208464e-02, -1.09267943e-01, + 1.43063843e-01, -6.33983240e-02, -1.62469558e-02, -1.03259414e-01, + -8.94914567e-02, -3.48297022e-02, -1.44557551e-01, 1.30005326e-04, + 1.44872010e-01, -1.74515136e-02, 4.52921912e-02, 1.08484700e-01, + -9.68603417e-03, -4.61588614e-03, -3.65241021e-02, -3.43756750e-04, + 1.57616019e-01, 8.44166130e-02, 2.49973107e-02, -9.32366401e-02, + 2.59006135e-02, 4.68501188e-02, 5.38958088e-02, 1.02374971e-01], dtype=float32), + kernel: Array([[-0.35964 , 0.58528113, 0.12339137, ..., -0.2728965 , + -0.35934824, -0.06503713], + [ 0.35246447, -0.02949085, 0.5403422 , ..., -0.09529549, + 0.11235058, -0.22131431], + [ 0.13894935, -0.3717892 , -0.66237545, ..., -0.41706035, + 0.52529556, 0.16840486], + ..., + [-0.09083151, 0.03931169, 0.2619683 , ..., -0.20559356, + 0.15253511, 0.3468645 ], + [-0.02631908, -0.0662327 , 0.12175545, ..., 0.10028415, + -0.3420218 , 0.01938576], + [-0.15720125, 0.16154028, -0.15429054, ..., -0.42136976, + 0.29129225, 0.06028228]], dtype=float32), + }, + }, + }, + }, + 2: { + attention: { + attention: { + key: { + bias: Array([ 4.91330713e-01, 9.67662875e-03, 5.60255200e-02, 8.97967637e-01, + -3.06834251e-01, 1.06595111e+00, 5.69995821e-01, -9.72940326e-01, + -1.00492425e-01, 2.51781523e-01, -3.73477489e-01, 1.62789464e-01, + -3.40229303e-01, 8.37876360e-05, 4.15999711e-01, 8.95448923e-02, + -1.16417743e-01, 1.02886930e-01, -1.22432733e+00, -2.46076301e-01, + -5.80214441e-01, -4.96384725e-02, -2.68343747e-01, -3.29451412e-01, + 2.24202454e-01, -6.01632036e-02, -6.19756997e-01, -3.03309470e-01, + -1.28122672e-01, -2.58336246e-01, 4.15848672e-01, 6.79572225e-02, + 1.68727383e-01, -1.08576640e-01, 9.50729072e-01, 1.18212245e-01, + 2.77229160e-01, 4.02590066e-01, 2.45219976e-01, -2.20667735e-01, + -6.55377328e-01, -2.22052753e-01, -2.34352529e-01, 7.21003652e-01, + -2.47502819e-01, -8.39031518e-01, 7.14816153e-01, 5.66149652e-01, + 1.96338549e-01, -3.21684837e-01, 1.26060262e-01, -1.13100424e-01, + 7.06214011e-02, -3.80029410e-01, -5.25692463e-01, 5.33015609e-01, + -2.28610292e-01, 8.77432287e-01, 4.54767972e-01, 1.05058357e-01, + -1.01594612e-01, -4.25737798e-01, 8.72221589e-01, -4.27799642e-01, + -7.26574242e-01, 3.97354126e-01, 5.40428281e-01, -6.92076981e-01, + -7.33269453e-01, -8.38262439e-01, -9.10226703e-01, -2.79593736e-01, + 5.23640335e-01, 2.84005404e-01, 3.39674622e-01, -2.09556028e-01, + -1.83250681e-02, 2.20996439e-01, 8.39932729e-03, -1.35142893e-01, + -4.47591990e-01, -5.44816703e-02, 3.56293529e-01, 3.84467632e-01, + -4.40692216e-01, 1.00856110e-01, 8.52475286e-01, 3.95925105e-01, + -5.06571472e-01, 4.37813222e-01, -8.72387230e-01, -4.80205119e-01, + -6.85112596e-01, -6.57971025e-01, 3.89101595e-01, -2.59548306e-01, + -2.97444463e-01, 2.12126691e-02, -8.14053118e-02, 1.06517822e-01, + 2.81973779e-01, 7.66788945e-02, -4.01952624e-01, -1.77236781e-01, + 6.38018996e-02, -3.04998636e-01, -3.08628559e-01, 1.56093314e-01, + 1.61672279e-01, 1.33270577e-01, 1.59121871e-01, -1.46263152e-01, + 1.54846618e-02, 2.81419247e-01, -3.13295014e-02, 2.87094891e-01, + -3.95765811e-01, 4.69570488e-01, -4.26991880e-01, 2.32855052e-01, + -1.62792206e-01, -3.81025262e-02, 4.74744802e-03, -7.95326605e-02, + 1.28198788e-02, 1.01442777e-01, -4.87101451e-02, -3.22961807e-02], dtype=float32), + kernel: Array([[-0.20145601, -0.22677487, -0.05235473, ..., -0.06058222, + -0.01446399, 0.16699584], + [ 0.22185628, 0.15144752, 0.23851915, ..., -0.10337625, + -0.07871008, 0.01390354], + [ 0.01399281, 0.06500424, -0.15749013, ..., 0.14727156, + -0.08356368, 0.03307934], + ..., + [ 0.01401845, -0.0230756 , -0.0378951 , ..., -0.40631384, + -0.10683019, 0.05925401], + [-0.00425362, -0.07440671, 0.03620811, ..., 0.21968801, + -0.18367934, 0.02989441], + [-0.1020968 , 0.5364823 , 0.00217419, ..., -0.05559423, + 0.13003536, -0.06085972]], dtype=float32), + }, + query: { + bias: Array([ 0.3528763 , 0.3471857 , 0.21516396, 0.44737092, -0.4285367 , + 0.3081475 , -0.25559098, -0.5126486 , 0.21324182, -0.32549337, + 0.02813848, -0.38898697, -0.34589097, -0.19787419, 0.28709203, + -0.24419545, -0.26928642, 0.30525193, -0.32151905, 0.0448105 , + 0.51535946, -0.11308531, 0.3212626 , 0.4922951 , 0.26803052, + -0.38731664, -0.2744102 , 0.39860168, 0.4643956 , 0.20418273, + 0.14548776, 0.0657926 , -0.21296258, 0.15750653, -0.3648503 , + 0.61559546, -0.35865372, 0.0193048 , 0.31175947, 0.2440965 , + 0.318281 , -0.7344339 , -0.88320374, -0.05781611, 0.18718946, + 0.20290868, -0.23830453, 0.5799299 , 0.01096278, -0.4863952 , + -0.03866521, 0.30802044, 0.44045374, -0.03443577, -0.02167755, + -0.01554048, -0.511046 , 0.41503018, -0.58751154, 0.5616692 , + -0.26840666, -0.20692529, 0.71333766, 0.2290065 , 0.15057702, + 0.1327177 , 0.31366843, 0.3856127 , -0.18231715, -0.24914798, + 0.321846 , -0.3574352 , -0.12077462, 0.36877826, 0.38871482, + -0.3660267 , 0.7790568 , 0.06110919, -0.22040306, -0.23536305, + -0.07361323, -0.47640693, 0.53327733, 0.49116883, -0.5596567 , + -0.7786464 , -0.1766858 , -0.14165151, -0.4430259 , 0.38748688, + 0.02364437, -0.46849003, -0.23776187, -0.45196384, 0.65864664, + 0.39194986, 0.07788196, 0.10206337, 0.15233228, 0.01135532, + 0.08919958, -0.0229385 , -0.00552564, -0.6363855 , 0.04719624, + 0.07559558, 0.0246234 , 0.32256266, -0.02410899, -0.37407702, + -0.063107 , -0.24087183, -0.0432428 , 0.01937803, -0.12409952, + 0.38090006, -0.8962724 , 0.51235485, -0.21019219, -0.06553419, + -0.0576391 , -0.00670896, 0.0239524 , 0.00629916, 0.46627957, + 0.02785698, 0.03592508, 0.06080098], dtype=float32), + kernel: Array([[ 0.09377078, 0.06110698, -0.01986536, ..., -0.09585967, + -0.00882398, -0.09038808], + [-0.1099542 , -0.0294884 , -0.3200316 , ..., 0.14632395, + -0.085163 , -0.12561767], + [-0.17357844, -0.08655021, -0.17681202, ..., 0.2461566 , + -0.11222455, 0.1451854 ], + ..., + [ 0.01483 , 0.12659054, 0.3020677 , ..., -0.1301052 , + -0.02523333, 0.0047435 ], + [-0.21112014, -0.10351861, 0.12283702, ..., 0.04758199, + -0.02335739, -0.20571615], + [ 0.04751823, -0.08418679, 0.021736 , ..., 0.23970562, + -0.10594929, -0.2381722 ]], dtype=float32), + }, + value: { + bias: Array([-0.20210104, -0.22905709, -0.14039448, -0.03683947, 0.12594987, + -0.02302479, 0.00906336, -0.00965133, -0.07539421, 0.19282112, + -0.14035326, -0.02629162, -0.0613457 , 0.19251512, 0.22437336, + 0.01865598, 0.09472503, 0.07302232, 0.23662573, 0.1634112 , + -0.25386894, 0.03644704, -0.14476314, -0.15108086, 0.30487335, + -0.01728077, -0.03989496, 0.10778196, 0.01612386, -0.02545331, + 0.10517047, 0.00122259, 0.06249418, 0.1109454 , 0.07476445, + -0.00254933, -0.08234914, 0.00984051, -0.18144692, 0.16224049, + 0.0791403 , 0.03824679, 0.05325351, 0.0459534 , -0.01707704, + -0.05914249, -0.07880631, 0.09950723, 0.01833864, -0.12028188, + 0.0201016 , 0.17571941, 0.08609623, -0.07171538, 0.27483281, + 0.19407238, 0.01660765, 0.03474303, 0.09089345, 0.19815543, + 0.21818227, 0.1436114 , 0.00318472, 0.06330492, -0.06290373, + 0.03819067, -0.06003238, -0.03846627, 0.07632727, -0.05140862, + -0.07411028, 0.21224949, -0.07434353, -0.09569198, -0.01517388, + 0.00378666, 0.0954181 , 0.02467517, -0.0480962 , -0.07227837, + 0.08442853, -0.11686811, 0.04613666, -0.04021059, 0.03575824, + 0.02964796, 0.09947117, 0.07083944, -0.06642204, -0.07676494, + 0.02942543, -0.0550985 , -0.06826791, 0.06714635, -0.05162644, + -0.03433613, -0.00067884, 0.05411651, -0.01232108, 0.10868998, + 0.05962406, -0.00136539, 0.01625578, -0.00632739, -0.0961687 , + -0.07040943, 0.02730714, 0.00544707, 0.01023827, 0.03284573, + 0.02555647, 0.0085297 , 0.03900068, 0.02747485, -0.03433363, + -0.01524182, -0.00976253, -0.06111204, 0.10030263, 0.04285727, + -0.00326952, -0.04434157, 0.11089424, 0.04277305, 0.02115745, + 0.02765508, 0.08399435, 0.08146379], dtype=float32), + kernel: Array([[-1.31283207e-02, 5.85139990e-02, 6.22775187e-05, ..., + 4.85775769e-02, -4.02368233e-02, -6.80026598e-03], + [-2.73164004e-01, -3.41370434e-01, -2.39238873e-01, ..., + 4.36691791e-01, 3.35174829e-01, 8.86783656e-03], + [ 6.34363815e-02, 3.22204083e-01, 2.63729244e-01, ..., + -9.32745263e-03, -1.88997626e-01, 2.92390492e-02], + ..., + [-1.07346885e-01, -1.71600416e-01, -1.54479533e-01, ..., + -2.66532987e-01, -7.03294575e-02, -3.63519728e-01], + [-1.77037895e-01, -1.31959140e-01, 3.59883085e-02, ..., + 1.12166613e-01, -3.25531811e-01, 4.15854841e-01], + [ 1.20714135e-01, 1.70483068e-01, 6.37787461e-01, ..., + -3.24400365e-01, -7.82848775e-01, -1.08967312e-01]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-1.42582014e-01, 3.51325125e-01, -5.49227633e-02, -1.36979753e-02, + -8.89001694e-03, -4.82381359e-02, 7.59436563e-02, 4.62903231e-02, + 1.25182196e-02, 1.74479246e-01, 5.47183445e-03, -4.21341211e-02, + -1.15092970e-01, 5.72893657e-02, -5.69992810e-02, 5.97054660e-02, + 2.31458675e-02, -9.98827890e-02, -1.63127463e-02, -5.68029238e-03, + 1.06309690e-01, -5.58492132e-02, -2.88810022e-02, -5.41146100e-02, + -1.25681654e-01, 1.55297220e-02, -1.05181850e-01, 2.84042470e-02, + 8.00593048e-02, 1.12994172e-01, 6.64846674e-02, -8.10639188e-02, + 6.96494877e-02, -7.10230693e-02, 2.48000082e-02, -6.40962422e-02, + 5.03854305e-02, 2.99889389e-02, 2.71755550e-02, -1.32961541e-01, + -9.42786504e-03, 6.68036118e-02, -7.92975008e-01, 2.92515382e-02, + -1.21656477e-01, 6.97143525e-02, 9.77309272e-02, 9.14023891e-02, + -1.06936611e-01, -4.60654832e-02, -3.59379910e-02, 6.73950538e-02, + -3.20891030e-02, 1.82818487e-01, -4.29846235e-02, 1.71033782e-03, + -2.07776949e-01, -8.03176984e-02, 1.26594633e-01, 1.53843045e-01, + 9.50956270e-02, -1.38820022e-01, -5.31725921e-02, 3.35701555e-02, + 5.83558269e-02, -1.29202589e-01, -4.29277681e-02, 5.83481155e-02, + 1.43310845e-01, 3.71531360e-02, -1.16031744e-01, 1.25720399e-03, + -3.66386138e-02, -5.68410717e-02, 8.81127954e-01, -3.54486227e-04, + -1.14806622e-01, 5.00139184e-02, -1.25719443e-01, -2.92788483e-02, + -5.60584605e-01, 1.90249607e-02, 7.62136094e-03, 3.29709910e-02, + -1.73091292e-02, 1.49499878e-01, 1.71786591e-01, 1.50196671e-01, + -4.66331327e-03, 7.90596679e-02, -6.07569050e-03, 1.27625940e-02, + -6.54113144e-02, -1.46687955e-01, -2.00468884e-03, -7.76169971e-02, + -5.61596686e-03, -2.51500830e-02, -5.23540825e-02, 1.42120197e-01, + 1.58276126e-01, -3.11181620e-02, 4.16391157e-02, -9.11307335e-02, + 1.01979673e-01, -1.40924305e-01, -2.65695825e-02, 3.79197043e-03, + -9.74692628e-02, -5.42117134e-02, -1.30008578e-01, 3.88463400e-02, + 1.63258657e-01, -1.50325231e-03, 3.45914699e-02, 1.69843420e-01, + 1.55994203e-02, -1.33964028e-02, -1.95804201e-02, 1.02751255e-02, + 1.84371233e-01, 7.10138083e-02, 5.72973862e-03, -1.13018960e-01, + -8.22554808e-03, 5.45389391e-02, 1.77288230e-03, 1.88416634e-02], dtype=float32), + kernel: Array([[ 0.15479796, 0.53571814, 0.51493394, ..., -0.80009085, + -0.05134518, 0.13081257], + [-0.00908494, 0.42273548, 0.710739 , ..., -0.3385195 , + -0.41763583, 0.50498915], + [ 0.21792167, 0.02829238, -0.15942883, ..., 0.14649649, + -0.3748477 , 0.5899068 ], + ..., + [ 0.69713473, -0.10411429, 0.20231156, ..., 0.13196032, + -0.09708058, -0.21109731], + [ 0.5224463 , -0.00511254, -0.9390049 , ..., 0.06903978, + 0.24796766, 0.5909974 ], + [-0.6673977 , 0.24624978, -0.7719052 , ..., -0.09019069, + -0.29529244, -0.12763341]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([-2.69522518e-02, 7.07723275e-02, -1.92383816e-03, -1.69017464e-01, + 6.28127530e-03, 1.35174885e-01, 1.43263070e-02, -1.49136512e-02, + 1.43651441e-02, 2.47675721e-02, 2.31170952e-01, -1.36354938e-02, + -5.00643514e-02, 4.64095622e-02, 7.30278119e-02, 2.58173402e-02, + -7.96374157e-02, -9.20170024e-02, -1.74471997e-02, 1.98375121e-01, + 2.97931284e-01, 1.71703979e-01, -1.74371853e-01, -8.47198628e-03, + -2.85758138e-01, 2.43080053e-02, -2.33928546e-01, 1.65304869e-01, + 1.08051397e-01, -3.64561914e-03, -4.06746566e-01, -2.43781097e-02, + -1.94358621e-02, 2.50640996e-02, 1.21675162e-02, 3.38391075e-03, + 4.25222097e-03, -1.16181098e-01, -6.87912405e-02, -1.43533750e-02, + 6.30778596e-02, 5.16471975e-02, 2.78629586e-02, -6.81551844e-02, + -2.16658652e-01, 1.19643852e-01, -1.77933127e-02, 4.19648476e-02, + -6.94157332e-02, 8.80410895e-02, 3.07484251e-03, -2.26102751e-02, + 2.71407124e-02, -1.37352645e-02, 1.24339305e-03, -1.10182655e-03, + -4.25179452e-02, 3.90366204e-02, -6.09224401e-02, 7.32898619e-03, + -5.39776146e-01, 8.10408741e-02, 1.43810920e-02, -1.46246385e-02, + 1.65752345e-03, -2.01134533e-01, 1.59140024e-02, 4.93869092e-03, + 4.72723007e-01, -1.37103930e-01, 1.31060192e-02, -4.23336178e-02, + 4.09131646e-02, 3.04093435e-02, 5.28761335e-02, 3.03985029e-01, + -2.57239580e-01, 1.94385331e-02, 2.98021808e-02, -3.17842066e-02, + -3.57343435e-01, 6.69872388e-02, -2.05376260e-02, 6.03192970e-02, + -2.77649582e-04, 4.66996022e-02, 2.09306199e-02, 1.26473635e-01, + -1.99992478e-01, 6.02058368e-03, -9.60124508e-02, -6.86570778e-02, + -2.69254763e-02, -1.16648778e-01, 6.80445060e-02, 8.54630917e-02, + 6.16764985e-02, -6.65853033e-03, -6.19218759e-02, -3.73522252e-01, + -1.32694215e-01, 4.03534174e-02, 5.64134538e-01, 5.32403849e-02, + -7.38940388e-02, -9.58362669e-02, -1.27810821e-01, 5.41193448e-02, + 1.32223934e-01, 4.17461758e-03, 7.13659152e-02, 3.29776555e-02, + 3.19321036e-01, -3.77637222e-02, 2.86701359e-02, 9.89994928e-02, + 5.24967327e-04, 6.47896230e-02, -9.30845514e-02, -8.97478983e-02, + 3.75827640e-01, 7.06133619e-03, -1.70939237e-01, -2.22015311e-03, + -8.06940794e-02, -1.80535465e-01, -1.17046207e-01, -5.05085103e-02], dtype=float32), + scale: Array([-2.16153208e-02, 4.99568492e-01, 1.90679863e-01, 3.98646802e-01, + 3.23043484e-03, 8.04649472e-01, 1.73618346e-01, 6.32202104e-02, + 6.03341032e-03, 7.60933161e-01, 5.64841628e-01, 3.30200821e-01, + -1.13521464e-01, 4.35642749e-01, 7.02156186e-01, 1.86426133e-01, + 2.69546956e-01, 4.07335073e-01, 3.06657463e-01, 7.59951055e-01, + 8.35915983e-01, 5.73645175e-01, 3.52393031e-01, 2.92183794e-02, + 5.95538437e-01, -3.08835153e-02, 5.68952084e-01, 6.18548512e-01, + 7.05874801e-01, -1.11141577e-01, 7.34429657e-01, 1.59030482e-01, + 6.69286311e-01, 3.07584852e-01, 2.97907770e-01, 2.45770905e-02, + 3.42297733e-01, 6.38300121e-01, 4.57696557e-01, 3.39488059e-01, + 4.27823365e-01, 6.61583364e-01, 2.51028594e-02, 6.87740266e-01, + 6.38814211e-01, 6.12671137e-01, 3.40550840e-01, 3.15654457e-01, + 5.78372836e-01, 5.64573705e-01, 5.55405932e-05, 3.71645719e-01, + 8.48703533e-02, 9.91551802e-02, 4.97666523e-02, 5.78866065e-01, + -1.33863851e-01, 2.80352771e-01, 5.00911117e-01, 1.85352847e-01, + 7.25611448e-01, 4.00964379e-01, 2.84439921e-01, 3.20623189e-01, + -1.70927703e-01, 4.05740440e-01, 3.79946142e-01, 2.70140201e-01, + 5.38314104e-01, 6.14893734e-01, 5.95924675e-01, 5.98044157e-01, + 1.48603588e-01, 3.81722450e-01, 5.22187948e-01, 7.07299054e-01, + 4.85696197e-01, 1.27313018e-01, 7.40394592e-02, 3.15272629e-01, + 3.43401968e-01, 6.31168902e-01, 2.14417100e-01, -3.38677377e-01, + 2.39464089e-01, 4.71792251e-01, 3.14170390e-01, 8.54388118e-01, + 6.38605714e-01, 1.02670453e-01, 4.03134346e-01, 3.42762381e-01, + 4.40316945e-01, 5.73590755e-01, 3.09652478e-01, 1.55366600e-01, + 6.81986809e-01, -7.89772999e-03, 5.20319641e-01, 6.24004126e-01, + 5.56326687e-01, 3.05760831e-01, 9.59146619e-01, 5.92343211e-01, + 3.85424703e-01, 4.82657462e-01, 5.25808513e-01, 4.87956852e-01, + 3.75198275e-01, -6.13166541e-02, 5.11129975e-01, 6.06624246e-01, + 7.52410650e-01, 1.93235308e-01, -5.42348363e-02, 3.70484799e-01, + 6.81148842e-03, 7.82989919e-01, 6.00687087e-01, 4.27850395e-01, + 9.20885563e-01, 2.06756860e-01, 7.61731803e-01, 6.10257149e-01, + 2.84399986e-01, 6.83825493e-01, 5.32033622e-01, 1.40386716e-01], dtype=float32), + }, + layernorm_before: { + bias: Array([ 1.7358573e-02, -2.1321216e-01, -3.0296935e-02, -8.5394399e-04, + -3.9229371e-02, 5.5817175e-03, -2.7252757e-03, -4.4029057e-02, + -2.0724675e-03, -1.2441806e-01, -2.2284750e-02, -2.7261129e-02, + -2.7341412e-03, 3.6942791e-03, 4.5538493e-03, 7.4563054e-03, + -6.2132953e-03, 9.7167321e-02, -2.5345799e-02, -3.9951671e-02, + -4.3273875e-03, -1.4203915e-03, 6.3152872e-02, 1.3384975e-03, + -5.5942591e-02, 6.3980678e-03, -3.2894940e-03, 2.6343842e-03, + 2.8435083e-02, 2.8260322e-02, 1.2178370e-02, -5.3195474e-03, + -2.4072591e-03, 4.1922098e-03, -7.3268535e-03, 7.5856894e-03, + 3.8502638e-03, -9.8595031e-02, 2.4951894e-02, -9.5122019e-03, + -5.4523430e-04, 4.9118581e-03, 2.0277210e-01, -1.6341466e-02, + 3.9525043e-02, -2.2266436e-02, -2.0717401e-03, -1.2440056e-01, + 2.4503253e-03, 1.1633219e-02, 3.5871702e-04, 1.0992530e-03, + 7.6659650e-02, -6.7069583e-02, 9.7138679e-04, -2.9395521e-02, + -8.4895557e-03, -1.3592694e-02, 2.9208256e-02, 1.7238649e-02, + 5.7587404e-02, 4.4692371e-02, 1.6647557e-02, -6.6133849e-03, + -3.3748765e-03, 9.0554915e-03, 1.8321996e-03, 1.5999395e-02, + -1.3479884e-02, 2.9679830e-03, -2.7784597e-02, 3.4300985e-03, + 4.0611736e-02, -4.6936791e-02, -7.2169447e-01, 8.0935545e-03, + -1.8306768e-02, -5.0263423e-02, -3.8130886e-03, 1.9901623e-03, + 2.5550249e-01, 6.1443527e-03, -2.2893531e-02, -1.7245837e-02, + 2.9439591e-02, -4.3969110e-02, -6.7670122e-03, -1.3685081e-02, + -9.9776837e-04, 9.0234309e-02, -1.4077160e-02, 2.6768299e-03, + -3.8174929e-03, -1.4334331e-02, 3.4874912e-02, 8.0644749e-02, + 8.6499676e-03, -1.7143000e-02, 7.9911202e-03, 2.8301086e-02, + -4.1625272e-02, -3.0587490e-03, -5.6571878e-02, 9.0596341e-03, + -7.6397270e-02, -4.5915369e-02, -6.2590851e-03, 3.4895688e-02, + 6.1382476e-02, -5.6433260e-02, 1.0131227e-02, 1.5290984e-03, + -5.1695153e-02, 1.3336526e-01, 1.1029670e-02, -2.2941723e-02, + 2.4411927e-03, -2.4650490e-02, -1.1010250e-02, -7.5843465e-03, + -1.1662335e-02, 3.2463686e-03, -1.6854508e-03, 2.7790260e-02, + 7.3739677e-03, -9.9853233e-02, 1.0033315e-02, -9.9916525e-02], dtype=float32), + scale: Array([ 7.94857554e-03, 1.53508723e-01, 7.40691870e-02, 1.35285795e-01, + 2.26536706e-01, 1.69742465e-01, 8.65431502e-04, 3.19571525e-01, + 1.62680820e-03, 1.62808627e-01, 1.34147704e-01, 2.01822311e-01, + 6.27691345e-03, 7.19762873e-03, 3.74700986e-02, 3.44181597e-01, + 9.95386466e-02, 1.60152465e-01, 9.15851742e-02, 1.63604781e-01, + 1.17830604e-01, -1.66135898e-03, 2.23148689e-01, 1.00941835e-02, + 2.60033756e-01, 1.75621884e-03, 2.65612965e-03, 4.38557938e-04, + 3.29755515e-01, 2.43666410e-01, 3.20717841e-01, 1.77067131e-01, + 2.68771589e-01, 3.02107055e-02, 2.95764953e-02, 1.08823786e-02, + -1.60506088e-02, 2.95251936e-01, -1.78114280e-01, 8.72005969e-02, + -8.84137396e-03, 2.99656868e-01, 8.63162801e-02, 1.33525252e-01, + 3.03551525e-01, 2.60979742e-01, -5.50631993e-03, 2.83051997e-01, + 2.56603450e-01, 2.00551033e-01, 5.14968205e-03, 2.78209685e-04, + 2.23627791e-01, 2.34901443e-01, -8.34921002e-03, 2.24900410e-01, + -6.23680511e-03, -2.86224540e-02, 2.08953783e-01, 3.78563493e-01, + 2.40491629e-01, 1.92855686e-01, 2.62306392e-01, 5.38155213e-02, + 2.10460991e-01, 1.67684779e-02, 5.48215546e-02, 1.18601710e-01, + 3.43158811e-01, 2.41343841e-01, 2.74241388e-01, -6.84597194e-02, + 2.30040535e-01, 2.19160318e-01, 1.71067789e-01, 2.74585098e-01, + 1.48560002e-01, 2.04053581e-01, -1.04316017e-02, 3.06529433e-01, + 1.51589945e-01, -2.79646995e-03, 2.60486454e-01, 9.33720767e-02, + 2.20337927e-01, 1.95606187e-01, -4.61113850e-05, 7.93705583e-02, + 2.31670171e-01, 1.81787178e-01, 6.02688640e-02, 6.66163536e-03, + 5.13362838e-03, 8.27364624e-02, 3.05262417e-01, 1.02641329e-01, + 1.57344952e-01, 1.41675845e-01, 4.87802178e-03, 1.00217178e-01, + 1.87011242e-01, -1.03036012e-03, 2.17317730e-01, 1.16062611e-01, + 3.36587638e-01, 3.04190874e-01, 1.32016614e-01, 1.66113153e-01, + 3.42909336e-01, 2.46614054e-01, 2.17617258e-01, 1.86562017e-01, + 1.04256548e-01, 3.53437752e-01, 5.86070295e-04, 2.16139719e-01, + 2.78897047e-01, 2.12332204e-01, 3.80921900e-01, 1.19425990e-01, + 3.23329829e-02, -1.03264116e-02, 1.89511199e-02, 1.55454144e-01, + -3.36376615e-02, 2.82643199e-01, 3.45727682e-01, 2.41249815e-01], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-2.87919283e-01, -1.59886956e-01, -3.92205477e-01, -4.32191879e-01, + -3.32138002e-01, -5.40607274e-01, -3.54391366e-01, -1.31854281e-01, + -5.39756715e-01, -4.69244629e-01, -3.44658494e-01, -4.88029271e-01, + -3.98942500e-01, -4.15098637e-01, -3.66716921e-01, -5.07466793e-01, + -2.69494671e-02, -3.29423159e-01, -3.49476814e-01, -1.71700358e-01, + -3.80648106e-01, -4.23726380e-01, -4.98886943e-01, -4.94191587e-01, + -2.60305732e-01, -3.48132432e-01, -4.27568406e-01, -3.35825920e-01, + -3.69937718e-01, -3.21786374e-01, -2.29871064e-01, -2.88744688e-01, + -3.85481507e-01, -4.31343943e-01, -1.41554534e-01, -1.79118499e-01, + -4.28618819e-01, -2.48103037e-01, -9.70930979e-02, -4.53777105e-01, + -1.97259694e-01, -3.89665425e-01, -4.33900356e-01, -2.01605186e-01, + -3.11053962e-01, -3.33365947e-01, -4.72573489e-01, -4.81180519e-01, + -4.59191829e-01, -3.63978118e-01, -5.11594951e-01, -3.21153313e-01, + -5.05568504e-01, -4.40596819e-01, -2.05569610e-01, -2.04613134e-01, + -2.30359152e-01, -4.58120763e-01, -1.35851100e-01, -2.78527498e-01, + -5.32928765e-01, -4.59086925e-01, -1.91062093e-01, -2.50361383e-01, + -3.23409230e-01, -5.09010673e-01, -2.47201502e-01, -3.10001701e-01, + -3.00568342e-01, -3.73322159e-01, -5.35207331e-01, -5.31656861e-01, + -3.18275809e-01, -4.40275580e-01, -3.43852162e-01, -4.12837327e-01, + -3.85347039e-01, -4.52745438e-01, -2.02980176e-01, -4.54181701e-01, + -5.00234187e-01, -3.34358454e-01, -3.39169592e-01, -2.89554566e-01, + -3.29827011e-01, -5.80846727e-01, -3.28373313e-01, -4.26805854e-01, + -3.15485120e-01, -3.92600238e-01, -3.12504858e-01, -2.19271436e-01, + -1.64767921e-01, -5.02310932e-01, -2.58719236e-01, -4.32749659e-01, + -9.22016874e-02, -4.70086277e-01, -3.97373617e-01, -4.89154935e-01, + -4.44214731e-01, -2.03317106e-01, -4.75122511e-01, -2.69819319e-01, + -2.67087281e-01, -4.90763247e-01, -4.22214806e-01, -2.64238060e-01, + -2.23214269e-01, -3.43690515e-01, -2.78915852e-01, -4.24799353e-01, + -5.69192111e-01, -2.07584277e-01, -1.64182380e-01, -1.91950992e-01, + -5.61932206e-01, -5.16008019e-01, -3.33474815e-01, -4.25828815e-01, + -1.61096185e-01, -3.98274213e-01, -5.68151116e-01, -2.51645744e-01, + -2.21852303e-01, -4.22558784e-01, -3.82831007e-01, -3.19712460e-01, + -1.66578472e-01, -2.49394879e-01, -4.19118315e-01, -2.93457031e-01, + -2.79008389e-01, -5.07268071e-01, -1.48156717e-01, -4.60653484e-01, + -2.05382407e-01, -4.70562130e-01, -2.83998400e-01, -4.29632753e-01, + -2.36523062e-01, -3.75578731e-01, -3.47366124e-01, -2.43162408e-01, + -4.38231677e-01, -5.36709070e-01, -1.51090562e-01, -4.71059680e-01, + -4.89868373e-01, -3.91524017e-01, -3.24053317e-01, -3.04224849e-01, + -4.36558515e-01, -2.35701188e-01, -4.07339513e-01, -2.95855463e-01, + -2.55902320e-01, -4.83437032e-01, -3.88129205e-01, -1.03971243e-01, + -2.60267794e-01, -3.92015398e-01, -4.49646682e-01, -2.48017311e-01, + -2.56850272e-01, -1.90488473e-01, -4.68013018e-01, -4.13734376e-01, + -5.12630165e-01, -4.26820993e-01, -3.74209106e-01, -3.59761298e-01, + -4.02847528e-01, -2.97854811e-01, -3.58279765e-01, -3.24418694e-01, + -3.66349548e-01, -4.33402061e-01, -4.91001993e-01, -5.03202379e-01, + -5.00787556e-01, -3.92910451e-01, -2.84466565e-01, -3.34319711e-01, + -4.25348043e-01, -2.26335734e-01, -2.37801537e-01, -4.87084776e-01, + -3.79190952e-01, -3.08594137e-01, -3.53578776e-01, -2.29536459e-01, + -4.85832751e-01, -3.31020027e-01, -3.49860758e-01, -4.98145282e-01, + -2.17634872e-01, -3.08456481e-01, -4.34388995e-01, -2.14153975e-01, + -2.20951349e-01, -4.75258082e-01, -2.87501752e-01, -9.52811018e-02, + -5.12147367e-01, -2.70634055e-01, -2.17771903e-01, -3.38476300e-01, + -3.06095690e-01, -3.33468020e-01, -1.24255188e-01, -5.16220808e-01, + -2.23946184e-01, -5.59429586e-01, -2.83430904e-01, -4.03296888e-01, + -3.68934721e-01, -3.25362712e-01, -7.39004537e-02, -3.09295535e-01, + -4.49250668e-01, -5.65693378e-01, -2.64678180e-01, -3.42923045e-01, + -3.37939441e-01, -3.71786445e-01, -4.06780869e-01, -5.13805926e-01, + -1.10545248e-01, -4.07462388e-01, -4.62072462e-01, -3.73685181e-01, + -4.82776046e-01, -1.83113769e-01, -1.37951612e-01, -4.32110876e-01, + -4.39051688e-01, -4.20039445e-01, -4.49939400e-01, -5.46136618e-01, + -4.43518162e-01, -3.36255163e-01, -2.60323614e-01, -1.48123205e-01, + -3.08105409e-01, -5.92286289e-01, -3.61411870e-01, -5.49996793e-01, + -4.72203493e-01, -1.12566963e-01, -3.46157908e-01, -2.21687108e-01, + -1.49518579e-01, -4.22669083e-01, -5.09057462e-01, -3.78608346e-01, + -2.69701719e-01, -2.57610738e-01, -3.43267351e-01, -3.99610460e-01, + -5.19208014e-01, -3.55502129e-01, -3.64218682e-01, -4.58064228e-01, + -1.66265354e-01, -4.56624329e-01, -4.21405643e-01, -3.38668913e-01, + -4.08919096e-01, -2.49855325e-01, -2.68893301e-01, 5.37731161e-04, + -4.11918223e-01, -2.53745079e-01, -2.37736523e-01, -3.77250522e-01, + -3.40394109e-01, -4.99163866e-01, -3.26041222e-01, -3.24631035e-02, + -4.25087065e-01, -3.99343103e-01, -4.45631206e-01, -2.19975799e-01, + -4.91138816e-01, -3.16317588e-01, -1.97857752e-01, -4.38183516e-01, + -2.38349304e-01, -3.12527388e-01, -4.46206242e-01, -1.97425947e-01, + -3.54068547e-01, -5.62409818e-01, -4.79922563e-01, -4.07248676e-01, + -5.08961380e-01, -3.35608751e-01, -4.69772518e-01, -4.64478105e-01, + -2.38058373e-01, -3.87834966e-01, -4.91571218e-01, -1.16468966e-01, + -1.88293502e-01, -2.19764918e-01, -2.41493329e-01, -3.08782160e-01, + -3.80097032e-01, -2.52258271e-01, -3.71472836e-01, -3.24205458e-01, + -4.05035138e-01, -5.40303230e-01, -1.73389211e-01, -3.20962131e-01, + -2.16032237e-01, -1.91938236e-01, -1.65382430e-01, -4.07554179e-01, + -3.40726435e-01, -2.95837790e-01, -9.07381624e-02, -1.67001665e-01, + -3.45026761e-01, -4.43528742e-01, -4.61021096e-01, -2.25696385e-01, + -2.09777251e-01, -5.53110957e-01, -3.23776424e-01, -3.86852056e-01, + -5.13776362e-01, -1.70138955e-01, -4.78795141e-01, -2.65858293e-01, + -1.22696631e-01, -5.30045271e-01, -3.90655428e-01, -2.64281482e-01, + -4.28456664e-01, -1.40250340e-01, -2.27036089e-01, -3.66672546e-01, + -4.42133546e-01, -4.25859839e-01, -2.67484695e-01, -4.62734908e-01, + -3.37051451e-01, -4.50581521e-01, -7.01922104e-02, -3.47592622e-01, + -3.91412824e-01, -4.24011290e-01, -3.32346976e-01, -3.72510850e-01, + -3.09896857e-01, -1.88529536e-01, -3.71884912e-01, -2.54814714e-01, + -3.63058448e-01, -4.75363225e-01, -3.06162030e-01, -4.65719730e-01, + -1.81163147e-01, -4.85125124e-01, -3.95848036e-01, -4.33863193e-01, + -3.04795027e-01, -4.90597248e-01, -2.44484976e-01, -3.44704390e-01, + -3.45088750e-01, -3.29458416e-01, -2.94500917e-01, -2.62573928e-01, + -1.69639900e-01, -4.71560776e-01, -2.42599428e-01, -3.69580209e-01, + -3.42656642e-01, -3.11950177e-01, -1.34305000e-01, -1.57368377e-01, + -3.91964465e-01, -3.12808394e-01, -2.37698436e-01, -3.69058728e-01, + -1.38768077e-01, -2.69422363e-02, -4.55560267e-01, -3.79797220e-01, + -4.36610907e-01, -2.92106032e-01, -3.98713797e-01, -4.69899833e-01, + -1.97569534e-01, -3.96044046e-01, -1.28238350e-01, -3.10444951e-01, + -4.51549977e-01, -2.95563221e-01, -4.12408620e-01, -4.34255183e-01, + -2.34395146e-01, -4.32429850e-01, -2.51709074e-01, -3.16113442e-01, + -2.24395722e-01, -4.12795454e-01, -3.96394074e-01, -3.76580715e-01, + -2.59395570e-01, -5.12954056e-01, -3.32104653e-01, -4.36192185e-01, + -5.13575554e-01, -1.39632359e-01, -4.51251060e-01, -4.95827556e-01, + -2.47558579e-01, -3.28758836e-01, -4.00682449e-01, -1.53779313e-01, + -2.39558816e-01, -3.46421301e-01, -2.07403153e-01, -4.59828466e-01, + -4.13055509e-01, -5.19209445e-01, -3.49062324e-01, -3.61741424e-01, + -5.46588659e-01, -3.25105667e-01, -2.64346927e-01, -3.62703383e-01, + -3.54973465e-01, -3.57288867e-01, -3.40226203e-01, -1.87480614e-01, + -3.84213597e-01, -3.67905647e-01, -2.21579894e-01, -4.32286561e-01, + -3.49667877e-01, -2.67124563e-01, -2.31143355e-01, -2.09680378e-01, + -4.44781452e-01, -1.10374503e-01, -1.90934509e-01, -4.75128025e-01, + -4.90421832e-01, -3.84051919e-01, -3.26278806e-01, -3.82725775e-01, + -2.96725929e-01, -4.35144246e-01, -3.04508001e-01, -4.23557281e-01, + -3.32650095e-01, -2.34854251e-01, -2.51995057e-01, -3.14376324e-01, + -4.71209139e-01, -5.88116467e-01, -3.59949678e-01, -3.27794611e-01, + -5.37792504e-01, -2.25097418e-01, -5.30355632e-01, -5.33735216e-01, + -2.61757433e-01, -2.64203936e-01, -3.20991725e-01, -3.19024026e-01, + -4.41864610e-01, -1.51468694e-01, -6.96782321e-02, -4.49556112e-01, + -3.41252834e-01, -6.03692472e-01, -4.20485705e-01, -4.03241307e-01, + -4.20258135e-01, -4.89987731e-01, -2.22824812e-01, -1.52097628e-01, + -3.01070571e-01, -3.31862330e-01, -2.91463912e-01, -2.47648209e-01, + -4.18553919e-01, 8.17964971e-02, -4.42326665e-01, -4.14036870e-01, + -4.23043162e-01, -4.24916655e-01, -3.22944313e-01, -3.46353531e-01, + -3.03149909e-01, -4.78416532e-01, -3.12508672e-01, -3.28645647e-01, + -3.10102880e-01, 6.61208481e-02, -3.10083032e-01, -4.34432387e-01, + -2.68977493e-01, -3.87750000e-01, -3.06754440e-01, -2.68470585e-01], dtype=float32), + kernel: Array([[-8.63173604e-02, 5.54475226e-02, -7.44883567e-02, ..., + -1.60854563e-01, -1.92300782e-01, 1.49715483e-01], + [-2.55667627e-01, -5.76032937e-01, -2.72577375e-01, ..., + -1.66639566e-01, -9.29499939e-02, -8.04038703e-01], + [ 1.05862938e-01, -7.12434500e-02, -3.51943791e-01, ..., + -4.56090540e-01, 6.42790049e-02, 2.83889264e-01], + ..., + [ 3.32350917e-02, -2.16388881e-01, -4.15698439e-01, ..., + -2.22950518e-01, 1.03157982e-01, 4.05906420e-03], + [ 1.34180099e-01, 2.93415725e-01, -3.85918528e-01, ..., + 5.76932847e-01, -1.18126400e-01, 1.22173645e-01], + [-3.73824616e-04, 1.94809213e-02, 2.37895355e-01, ..., + -2.55678475e-01, -4.13769521e-02, 8.79469365e-02]], dtype=float32), + }, + output: { + bias: Array([-0.14310208, 0.30402932, -0.04476445, 0.05614336, -0.00636329, + -0.10514647, 0.06905492, 0.06197605, 0.02608861, 0.15183868, + -0.0424929 , -0.01462859, -0.1350344 , 0.02212298, -0.05124652, + 0.06123491, 0.03247906, -0.09240599, -0.01013029, -0.03223318, + -0.0887438 , -0.06346248, 0.02817548, -0.05015382, -0.05139951, + 0.01631817, -0.04831094, -0.02007524, 0.04097073, 0.1089523 , + 0.14585266, -0.05194474, 0.04915197, -0.09083902, 0.01939188, + -0.06052509, 0.05472679, 0.04450262, 0.03008708, -0.1406127 , + -0.0286117 , 0.04946226, -0.8201933 , 0.05301183, -0.06421784, + 0.03291068, 0.08303672, 0.10530804, -0.05531342, -0.09611814, + -0.02686962, 0.10026816, -0.01932072, 0.17867358, -0.03787487, + 0.01405966, -0.20247376, -0.07591344, 0.15851738, 0.16269532, + 0.17163317, -0.16734251, -0.05690756, 0.02852942, 0.06089685, + -0.04197973, -0.06455492, 0.07540473, -0.06004319, 0.08624789, + -0.1040978 , 0.00247931, -0.02840352, -0.04520692, 0.8720378 , + -0.06833719, -0.07094937, 0.05479927, -0.12080298, -0.0121475 , + -0.5797822 , -0.02564085, 0.03635592, 0.03610236, -0.00385189, + 0.1565094 , 0.14343542, 0.05017952, 0.00714095, 0.09029737, + 0.0294838 , 0.01819415, -0.04793838, -0.1418107 , -0.00120118, + -0.06547695, -0.03586074, -0.01915462, -0.03162472, 0.25057518, + 0.21378383, -0.03219284, -0.03374225, -0.07630365, 0.13850403, + -0.09364668, 0.02092841, -0.02086668, -0.11372983, -0.06087056, + -0.14589557, 0.03926003, 0.0446431 , 0.00222743, 0.03577771, + 0.14117669, 0.01693195, -0.06432896, 0.02443499, 0.04788267, + -0.13781525, 0.06813591, 0.10359713, -0.10075218, 0.02409258, + 0.05323493, 0.00397177, 0.03774352], dtype=float32), + kernel: Array([[ 0.16250458, -0.15818997, 0.3543968 , ..., -0.06277009, + -0.458476 , 0.21236774], + [ 0.17157121, -0.33716652, 0.17193943, ..., -0.07994235, + 0.03173731, -0.09947567], + [ 0.20364675, 0.04469377, -0.3887421 , ..., -0.24748199, + 0.29170632, -0.02271716], + ..., + [ 0.29266545, -0.02244417, 0.41813877, ..., -0.37108174, + -0.26156428, -0.10771853], + [-0.07057951, -0.23378706, 0.05101302, ..., 0.08562624, + 0.45771068, 0.08117881], + [-0.21127233, -0.6031122 , -0.4248507 , ..., 0.16734137, + 0.09727449, -0.12857328]], dtype=float32), + }, + }, + }, + }, + 3: { + attention: { + attention: { + key: { + bias: Array([-0.22031489, 0.02730584, 0.23429762, -0.3125938 , 0.22592317, + -0.6771136 , -0.21247995, -0.341136 , 0.07302368, 0.44729188, + 0.22906882, -0.524169 , -0.04314182, -0.22182061, -0.67601067, + -0.01997758, -0.00890862, 0.31220517, 0.3243842 , -0.08772013, + 0.18925034, -0.43556792, -0.30325386, 0.13643748, 0.27908477, + 0.26076847, 0.19125256, 0.6231867 , 0.34878573, 0.21458477, + 0.12557186, 0.25881395, 0.11842591, 0.14228185, -0.14626004, + 0.02273282, 0.04946363, 0.3702873 , 0.4587399 , 0.46696767, + 0.2649066 , -0.02744056, 0.21641424, 0.4786166 , -0.39902607, + -0.25150967, 0.20756108, -0.9501807 , 0.2439707 , 0.07372095, + 0.1429905 , 0.5078466 , -0.47601166, 0.01509763, -0.63272023, + 0.5627352 , 0.5814683 , 0.3356227 , -0.07830408, -0.29685998, + 0.09222841, -0.49363494, -0.09527437, -0.01283222, 0.10691621, + 0.44450942, 0.2378995 , 0.14754085, 0.07311877, -0.4130656 , + 0.12185058, 0.1367209 , 0.18846467, -0.11077244, -0.10472129, + 0.13243866, -0.6801274 , 0.18685853, -0.40455377, -0.17242797, + -0.3273368 , 0.21752419, 0.5587714 , 0.33085388, -0.13658075, + 0.11016981, -0.14520068, -0.0551539 , -0.0988646 , 0.19754803, + 0.28224358, 0.27123928, 0.01996558, 0.25261068, 0.3291632 , + 0.73044914, -0.7135282 , -0.2712479 , 0.17239682, 0.4813283 , + 0.350012 , 0.76145405, 0.01346729, 0.28428316, 0.17260526, + -0.40829295, 0.66129416, -0.35163417, -0.15398747, -0.87850165, + 0.16401769, -0.36521718, 0.08581091, 0.12348412, -0.731215 , + 0.24665482, 0.77237046, 0.39620268, -0.4511514 , 0.34254473, + -0.669006 , -0.50109476, -0.18661116, 0.2477596 , 0.65949994, + 0.01988118, -0.04483053, -0.7621275 ], dtype=float32), + kernel: Array([[ 0.14074115, 0.11149215, 0.02833093, ..., 0.24423462, + -0.3480132 , -0.06901791], + [ 0.13614887, -0.2725882 , -0.31767422, ..., -0.30451745, + -0.0344476 , 0.2648898 ], + [ 0.19196813, 0.24347357, 0.09369045, ..., 0.3346163 , + -0.10084 , -0.19047955], + ..., + [ 0.02364591, 0.3210728 , 0.15941176, ..., -0.5476011 , + 0.30842453, 0.16627334], + [-0.3327589 , 0.2611727 , 0.04791668, ..., 0.5718048 , + -0.00553066, -0.12860055], + [ 0.13814813, -0.2723826 , -0.10484312, ..., 0.10856424, + -0.46137026, 0.3135762 ]], dtype=float32), + }, + query: { + bias: Array([ 0.01636963, 0.00303602, -0.2237812 , 0.17135 , -0.5504063 , + 0.71538967, 0.34468257, 0.6879721 , -0.32814807, 0.12084308, + -0.23467594, 0.24040897, -0.03144267, 0.41499555, 0.18407044, + -0.19885245, 0.29428306, -0.17422631, -0.72670996, 0.19728792, + -0.12558015, -0.04110402, 0.7530146 , -0.32672936, -0.34328216, + 0.39724907, -0.19968817, -0.0735909 , -0.11611477, 0.2887974 , + 0.27589712, -0.19556114, 0.00310108, 0.05990869, -0.06629719, + 0.11465285, 0.30233756, -0.4840145 , 0.36792257, 0.01198539, + 0.17746153, 0.00102492, -0.73527193, -0.04222427, 0.1278375 , + -0.04002868, -0.04392892, 0.07372226, 0.00746391, 0.01089494, + -0.06465401, -0.01894091, -0.04520243, -0.05135101, 0.10821903, + -0.36141765, 0.01786461, -0.00758973, 0.20370121, 0.07571532, + -0.04988701, -0.50066465, -0.29137757, 0.0468887 , -0.0369802 , + -0.28495568, -0.15369293, -0.04906723, -0.05060489, 0.23548573, + 0.02296117, 0.29195184, -0.4791245 , -0.32274112, 0.00163329, + 0.6024667 , -0.02721317, 0.4371868 , -0.45318347, 0.15298416, + 0.53022355, 0.23577152, -0.11320106, -0.05042947, 0.00792542, + 0.34630367, 0.45112005, -0.17176089, -0.30412972, -0.00977058, + 0.3797044 , -0.55427474, -0.14912651, -0.31753495, -0.14334434, + -0.12312431, 0.29336822, 0.08129798, -0.09826545, -0.37803188, + -0.03996738, -0.565557 , -0.15023306, -0.17836387, -0.00129342, + 0.31617662, -0.40511948, -0.0246349 , 0.8049514 , 0.50346047, + -0.08232877, 0.7958381 , 0.11525708, -0.87054974, 0.02044921, + 0.10254359, -0.41025582, -0.11836894, 0.3934775 , -0.2755682 , + 0.28480217, 0.28756127, -0.06242 , -0.36147273, -0.31789255, + -0.1169754 , 0.13862565, 0.5221684 ], dtype=float32), + kernel: Array([[-0.03366861, 0.1951855 , 0.13829306, ..., 0.20226802, + 0.14245163, 0.0851718 ], + [-0.12275442, 0.09568106, 0.1026165 , ..., -0.11756799, + -0.03399425, 0.1083973 ], + [-0.11897923, 0.44871688, 0.22821409, ..., -0.19063285, + 0.14112021, 0.05127604], + ..., + [ 0.42162004, -0.30832654, -0.31275368, ..., -0.48211238, + 0.221526 , 0.17857978], + [ 0.23314169, 0.38464716, 0.37320703, ..., 0.7464507 , + 0.65437895, 0.09728524], + [-0.02876819, -0.11770961, -0.18619719, ..., -0.4305184 , + -0.24435337, -0.0622173 ]], dtype=float32), + }, + value: { + bias: Array([ 3.80855761e-02, 1.19626271e-02, 9.54822153e-02, 3.39327864e-02, + 6.47817999e-02, 1.12312162e-04, 2.36410853e-02, 1.12216197e-01, + 1.71483278e-01, -4.48498651e-02, -6.95201615e-03, -5.79331703e-02, + 1.67439714e-01, -1.31781459e-01, 1.07828841e-01, -1.07748993e-01, + -2.16874084e-03, -2.65070498e-02, -1.59811676e-02, 7.09029958e-02, + 5.01474105e-02, -4.89775911e-02, 2.54220992e-01, -4.80323769e-02, + -1.79449767e-02, -5.27789593e-02, -4.24154028e-02, -1.70985550e-01, + -1.36494562e-01, 1.54837176e-01, -1.15875848e-01, -1.71242371e-01, + 5.86349219e-02, 1.91222817e-01, -2.35057212e-02, -8.43578875e-02, + -4.65269759e-02, -8.38292018e-03, -4.65860590e-02, -5.02779474e-03, + 4.50352840e-02, -2.31239274e-02, 5.93480729e-02, 5.71514890e-02, + 5.99062406e-02, 2.83261593e-02, -4.82074209e-02, -4.80451956e-02, + -6.26677051e-02, 5.32909296e-02, 6.67422032e-03, 4.72835749e-02, + 4.04058099e-02, -4.40718755e-02, -4.50864471e-02, 1.62461493e-02, + -5.04599810e-02, -3.60567644e-02, 4.91038598e-02, -7.88058713e-03, + 1.09715499e-01, 5.14943898e-02, 3.80876176e-02, 5.21656685e-02, + 1.56386979e-02, 6.83846846e-02, 7.60920197e-02, -8.33767503e-02, + -1.77296321e-03, -9.03388578e-03, -5.15193716e-02, 2.69258358e-02, + 4.77294379e-04, -2.96406243e-02, -8.59358441e-03, -2.96034720e-02, + -2.67447811e-02, -6.73482614e-03, 1.27622047e-02, -2.17594709e-02, + 8.59617069e-02, 1.03174895e-01, 2.63699442e-02, -1.10600488e-02, + -7.31578097e-02, -3.59972566e-02, 4.40209918e-02, -4.44998406e-02, + -1.85169447e-02, -4.11919132e-02, -3.58521268e-02, 1.47900637e-03, + 4.30759192e-02, 2.79228836e-02, 7.82085881e-02, 3.90527584e-02, + 3.11766323e-02, 2.85879765e-02, -8.20098892e-02, -7.59413168e-02, + -6.41580895e-02, 2.49105208e-02, -1.47475805e-02, -9.27800685e-02, + -1.73971094e-02, 8.51288959e-02, -2.54925918e-02, 4.36546579e-02, + -9.24658105e-02, -1.01948939e-02, -4.81284559e-02, -3.62936854e-02, + 2.90373899e-02, -4.27892022e-02, 3.08775119e-02, 5.98094938e-03, + 2.10134909e-02, 4.55508381e-02, 3.65342051e-02, 1.74732562e-02, + 1.06893666e-02, -1.42213493e-03, -2.61030179e-02, 2.22432707e-02, + 1.47231504e-01, 6.33382201e-02, 2.16640271e-02, 4.13991958e-02], dtype=float32), + kernel: Array([[ 0.00939619, 0.02482712, 0.01164271, ..., 0.0164959 , + 0.00769011, -0.03742562], + [ 0.18655218, -0.36816835, -0.08651444, ..., 0.2512268 , + 0.02853749, 0.00129432], + [ 0.09724977, -0.17311321, -0.15452595, ..., -0.4560767 , + -0.48346895, -0.4558498 ], + ..., + [-0.06453069, 0.4239059 , 0.45162615, ..., -0.27700382, + 0.13836603, -0.33547392], + [-0.02759907, -0.9605103 , -0.82787424, ..., -0.51210946, + 0.33198425, 1.0528204 ], + [-0.12098552, -0.10120461, -0.10430597, ..., 0.78033483, + -0.08541391, 0.05855659]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-1.72631651e-01, 2.81819075e-01, -1.31835248e-02, -6.50164206e-03, + -9.19895153e-03, -6.00125082e-02, 4.56261076e-02, -1.17319543e-03, + 3.18587162e-02, 2.30925724e-01, -2.78387107e-02, -7.64354365e-03, + -7.98858032e-02, 2.95766089e-02, -4.09963764e-02, 1.03628486e-01, + 5.56401908e-02, -8.70198682e-02, -1.29737146e-02, 1.11193627e-01, + -5.96015491e-02, -5.83102107e-02, -2.50512566e-02, -3.02468054e-02, + -3.29878554e-02, 4.26750369e-02, -3.15967947e-02, 1.79770100e-03, + 6.93126023e-02, 8.52573961e-02, 4.88307029e-02, -6.38186038e-02, + 1.12157442e-01, -9.84508917e-02, 5.63154779e-02, -7.03682825e-02, + 2.64801327e-02, -1.89797301e-02, 6.20345166e-03, -8.26654360e-02, + 2.39220913e-02, -9.58621129e-03, -8.39129508e-01, 4.56922092e-02, + -1.53480023e-01, 6.83549270e-02, 7.20679983e-02, 2.08752692e-01, + 4.47404385e-02, -5.50162159e-02, -5.08189276e-02, 5.22453226e-02, + -2.34786086e-02, 1.55734494e-01, -4.19281684e-02, -1.19545199e-02, + -8.40202495e-02, -8.96746367e-02, 6.41859993e-02, 7.71264508e-02, + 8.69383588e-02, -1.31521776e-01, -2.94522438e-02, 4.01933566e-02, + -1.60081685e-03, -3.63952629e-02, -5.23505732e-02, 8.07748511e-02, + 3.52987312e-02, 8.45032781e-02, -1.18315943e-01, 3.48625183e-02, + 5.90237789e-02, -7.91934412e-03, 8.07135582e-01, -6.57365983e-03, + -6.89911172e-02, 1.08262442e-01, -1.27490476e-01, 1.97968110e-02, + -6.29805326e-01, 5.13264053e-02, 3.17039192e-02, 1.63930492e-03, + 2.53316425e-02, 1.45042747e-01, 1.39836982e-01, 1.34169623e-01, + 4.38121334e-02, -1.00197628e-01, 3.51289548e-02, 3.51637453e-02, + -1.21273242e-01, -1.04602434e-01, -9.42274835e-03, -2.06037387e-02, + -4.44292687e-02, -2.39612963e-02, -1.44756138e-02, 9.55058485e-02, + 3.27826925e-02, 5.12430328e-04, 1.59999877e-02, -6.98894709e-02, + 1.01640105e-01, -5.48832715e-02, 6.11420302e-03, -7.91232064e-02, + -1.24512471e-01, -3.10322326e-02, -1.39527261e-01, 3.97435911e-02, + 1.18619666e-01, -1.12082511e-01, 1.25646042e-02, 1.72376752e-01, + 4.17521298e-02, -2.35856101e-02, -3.09479032e-02, 6.41909391e-02, + -2.64365948e-03, 2.64008511e-02, 1.21359736e-01, -4.78715599e-02, + 4.74941581e-02, 5.94909228e-02, -6.63625225e-02, -2.96208411e-02], dtype=float32), + kernel: Array([[ 0.30733457, 0.11394037, -0.11875194, ..., -0.02600583, + 0.26162902, -0.24450323], + [ 0.32579228, -0.01019196, 0.08372336, ..., -0.23464818, + 0.77629703, -0.14200607], + [ 0.40005246, -0.09073471, 0.01217228, ..., 0.12469898, + 0.4625953 , -0.2634586 ], + ..., + [-0.14610979, 0.26803997, -0.26479524, ..., -0.581134 , + 0.39575958, -1.1754091 ], + [-0.10739683, -0.4411798 , 0.68852794, ..., -0.01546267, + 0.06261325, -0.04542963], + [ 0.20287202, 0.13014646, 1.1704878 , ..., -0.26887265, + -0.26802245, 0.629474 ]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([ 1.71948429e-02, -1.09049575e-02, 6.19933335e-03, -3.97290401e-02, + -1.75763867e-04, 5.07760718e-02, 1.12305544e-02, -2.87154932e-02, + -2.46254611e-03, -1.75729260e-01, 9.95321851e-03, -1.72805071e-01, + -1.28897009e-02, 2.74409372e-02, -1.33280344e-02, 1.57673303e-02, + -2.70904624e-03, 1.06892176e-03, 2.89133284e-03, 1.36673786e-02, + 4.77137081e-02, -9.59468889e-06, 6.27165893e-03, -6.68704743e-03, + 1.20838732e-02, 1.39770750e-02, -3.87260579e-02, -5.88312745e-03, + 4.62619169e-03, 6.23012260e-02, 8.07429329e-02, 7.94423558e-03, + -2.50669499e-03, 1.25099311e-03, -6.49371138e-03, -9.59886145e-03, + 3.83810289e-02, 7.55729433e-03, 4.41730954e-03, 1.75167769e-02, + 3.34830210e-02, 9.28634033e-03, 1.21751567e-02, -1.52964927e-02, + -2.70186807e-04, 7.17550702e-03, 8.44905898e-02, -9.36958939e-02, + -1.57131581e-06, 4.74082539e-03, -8.86013955e-02, 7.37810042e-03, + -4.76047629e-03, 6.89025401e-05, -9.31238383e-03, 6.64415536e-03, + 7.76583795e-03, 2.29338557e-02, -1.50072891e-02, -3.97083163e-03, + 5.64871766e-02, 1.86805464e-02, 5.85068855e-03, 3.33538577e-02, + 3.15117603e-03, 3.10075725e-03, -3.19483913e-02, 3.20728831e-02, + 2.38816179e-02, 1.34205222e-02, -1.76607758e-01, -5.26308222e-03, + 2.82549043e-03, -1.20671205e-01, -1.07251152e-01, 3.50134517e-03, + 2.33329542e-04, -1.22928575e-01, -5.50287850e-02, -2.40997149e-04, + 4.36784141e-03, -3.80961671e-02, -1.24475360e-01, -1.74865853e-02, + 9.07492265e-02, -1.61794608e-03, 4.57073562e-03, -6.20833039e-03, + -5.41980378e-03, 2.80724547e-04, -4.64204960e-02, -4.63194735e-02, + 1.02222068e-02, -5.72193079e-02, -3.21766338e-03, 6.88658096e-03, + -4.31186426e-03, -3.94533984e-02, -1.81648252e-03, 5.71832731e-02, + -6.70271181e-03, 5.93083352e-03, -8.59189108e-02, -5.01573570e-02, + -3.59455287e-03, -8.11670348e-02, -3.07016950e-02, 2.01600092e-03, + 4.20039892e-02, -4.99186059e-03, -1.39956048e-03, -5.72945038e-03, + 1.23080192e-03, 1.21427193e-01, 1.27224373e-02, 5.32536209e-02, + -1.02717979e-02, -3.96808702e-03, 3.51909577e-04, 1.07945967e-02, + -4.73669693e-02, -5.23439702e-03, -3.12553300e-03, 6.25749454e-02, + -9.56553966e-03, -9.94754490e-03, 2.23868527e-02, -2.48279348e-02], dtype=float32), + scale: Array([ 4.24516723e-02, 1.21176187e-02, -4.84997500e-03, 5.57146311e-01, + 2.28432156e-02, 6.31359816e-01, -4.46775509e-03, 6.18052892e-02, + -4.33097221e-03, 2.61805564e-01, -1.88892316e-02, 6.95380867e-01, + -6.15348248e-03, 5.73870949e-02, 1.55512363e-01, -5.18870428e-02, + 1.60957292e-01, -1.62363220e-02, 4.54635024e-01, -1.21937275e-01, + 4.06006753e-01, 4.14683282e-01, 3.15475985e-02, 4.67635952e-02, + 1.88465670e-01, -1.32318744e-02, 4.83835369e-01, 8.42766091e-02, + -7.77654797e-02, 5.02031922e-01, 4.80292380e-01, 3.83823998e-02, + -4.09559421e-02, 1.13746740e-01, -1.26264850e-02, 7.38768047e-03, + 3.53148162e-01, -1.56045742e-02, 2.96747178e-01, 1.84800133e-01, + 2.40674987e-01, -5.23572266e-02, 1.17433036e-03, 6.38127849e-02, + 2.33978732e-03, -3.27690654e-02, 4.07037109e-01, 1.53259337e-01, + 3.08140907e-02, -5.92192961e-03, 3.31769854e-01, 3.03894043e-01, + -1.66736450e-02, -4.60295379e-02, 3.50147933e-01, 4.64291312e-02, + 7.68471509e-03, 9.69270691e-02, 1.85923763e-02, -9.66442749e-03, + 4.18877095e-01, 1.00185797e-01, 1.32984370e-01, 4.85280842e-01, + -7.05129327e-03, 4.97093098e-03, 3.28959554e-01, 9.48822647e-02, + 2.17689544e-01, 2.05316558e-01, 5.99317133e-01, 3.90748940e-02, + 5.26148966e-03, 3.09464663e-01, 3.66784409e-02, 4.52369526e-02, + -4.48510870e-02, 2.84545600e-01, 3.67504448e-01, 4.94225929e-03, + -5.27279684e-03, 8.71935904e-01, 4.51588362e-01, 2.58467104e-02, + 3.65057826e-01, 6.95686322e-03, -1.69470385e-02, 1.21857105e-02, + -1.65910535e-02, -8.54126003e-04, 5.98546147e-01, 3.60920489e-01, + 2.29799710e-02, 3.24622720e-01, -8.44317749e-02, -2.27606995e-03, + 3.23808223e-01, 2.70338178e-01, 5.69886751e-02, 1.97269484e-01, + 3.14833894e-02, -3.49436514e-02, 3.14999282e-01, 5.10905981e-01, + 2.21107565e-02, 3.05217803e-01, 2.38849521e-01, -1.36028053e-02, + 4.52475280e-01, 2.09674463e-02, -2.62735710e-02, -6.97454140e-02, + -8.80253036e-03, 4.10522908e-01, -2.87996326e-03, 2.52809376e-01, + -1.01799294e-02, 1.55074904e-02, 4.01195139e-02, 1.15418859e-01, + 1.29438698e-01, -9.92912054e-03, 3.01613122e-01, 9.44039226e-02, + -3.23632099e-02, 1.09247090e-02, 5.09600282e-01, 6.76498190e-02], dtype=float32), + }, + layernorm_before: { + bias: Array([ 2.0053135e-02, -2.0385702e-01, -5.2851941e-02, -1.5826782e-02, + -3.9745602e-03, -5.1497813e-02, 4.4322563e-03, -2.0497143e-02, + 3.5912704e-02, -1.5449695e-01, -6.0098185e-03, 4.4156699e-03, + -5.7948790e-03, 1.7218473e-03, 8.6186066e-02, -8.2310978e-03, + 8.4572081e-03, 3.2470144e-02, -7.2876806e-04, -7.4659437e-02, + -1.4426826e-02, -6.7664213e-03, 1.5026542e-02, 4.0727290e-03, + 2.3569600e-04, -9.6417527e-04, -1.1503830e-03, -1.6481975e-02, + -7.0209317e-03, 5.3436416e-03, 6.8159692e-02, -2.5559873e-03, + -8.0474456e-05, -3.0456469e-03, -3.3534174e-03, 6.1106024e-04, + -6.6513027e-04, -1.2051529e-01, -4.2777561e-02, 6.5030590e-02, + 2.1362038e-02, 3.2946058e-02, 2.9181463e-01, 1.6087948e-03, + 8.0183737e-02, 6.0054217e-04, -1.5031694e-03, -1.3727798e-01, + 2.0673402e-02, 1.6634881e-02, -6.1021522e-02, -9.7963065e-03, + -4.4847801e-03, -1.6810349e-03, -6.0316608e-03, -1.3532469e-04, + 8.6372629e-02, -1.0683674e-02, -1.5961612e-02, 4.4840686e-02, + 5.4682657e-02, 1.3794379e-02, 2.7690900e-03, 3.4621611e-02, + 3.8512382e-03, 1.3362155e-03, 2.2860983e-02, 1.7252501e-02, + 7.1021970e-03, -1.8353600e-02, -9.1281412e-03, -2.5052501e-02, + 6.4535923e-02, -1.1188607e-01, -7.4342960e-01, 9.8237805e-03, + -4.8992522e-02, -3.0361801e-02, 4.7444657e-05, -1.2857079e-02, + -3.5183441e-03, -2.8601723e-02, -4.8774560e-03, -4.9416926e-02, + 2.1382656e-03, -1.7278522e-02, 4.1791960e-03, -6.8051869e-04, + 5.9474696e-04, 1.3638437e-01, -9.9423109e-03, 1.4076917e-03, + -2.9004067e-03, 1.2549493e-02, 1.8330464e-02, 9.2371933e-02, + -8.4709952e-04, -1.5642725e-02, 1.8227950e-02, 4.0236771e-02, + -4.6799608e-02, -6.4027351e-03, -1.0166033e-01, 2.3036268e-02, + -2.3710703e-02, -2.6576657e-02, 4.9998323e-03, 2.3283776e-02, + 6.9294289e-02, -2.0691980e-02, 4.1049831e-02, 1.8761930e-04, + -1.0511643e-01, 8.2109414e-02, -1.6860457e-02, -3.8362190e-02, + 4.8856564e-02, -2.3737226e-02, -1.6952256e-02, 1.5742332e-02, + -1.9643296e-02, 7.6947906e-03, 1.0548260e-03, 5.7553727e-02, + 1.7045023e-02, -3.2199997e-02, 7.6176226e-02, -7.1644776e-02], dtype=float32), + scale: Array([ 1.36227664e-02, 1.40898958e-01, 1.92645520e-01, 4.70757820e-02, + -2.37463228e-03, 3.57880324e-01, 1.77441984e-01, 2.04336628e-01, + 5.04735708e-02, 2.16793969e-01, 1.84319422e-01, 1.62005424e-01, + -2.82178313e-04, 3.39963939e-03, 3.45121980e-01, 3.06478173e-01, + 4.18081768e-02, 7.82519281e-02, -1.85935143e-02, 2.79815465e-01, + 2.39701569e-01, 2.40149871e-02, 9.43208337e-02, 1.05237709e-02, + 9.96645987e-02, 5.39937872e-04, -1.31208654e-02, 1.39623404e-01, + -3.63838188e-02, 3.99398878e-02, 3.38648826e-01, 9.40667465e-02, + 1.93201587e-01, -4.52386551e-02, 3.65994535e-02, 3.85606177e-02, + -3.78566347e-02, 2.66506523e-01, 2.40695551e-01, 3.16709459e-01, + 8.35959762e-02, 2.61536777e-01, 1.21971466e-01, 1.21512767e-02, + 3.52942318e-01, 2.85213944e-02, -1.15962522e-02, 2.87642092e-01, + 2.67458588e-01, 3.02993476e-01, 1.73773259e-01, 1.12724066e-01, + -8.12414661e-03, 2.45114770e-02, 1.12504624e-01, 2.88635880e-01, + 1.51033133e-01, -3.74849811e-02, 3.15434486e-01, 3.69797438e-01, + 3.93722087e-01, 1.04381077e-01, 1.80823863e-01, 2.28892982e-01, + 2.66926765e-01, -5.73799293e-03, 1.90365806e-01, 2.51159281e-01, + 2.34797284e-01, 2.54733443e-01, 1.83505401e-01, 1.44614145e-01, + 2.63359755e-01, 4.59246188e-01, 2.01807499e-01, 2.74832934e-01, + 2.06239209e-01, 1.09179243e-01, 6.59588631e-03, 2.21557349e-01, + -5.72106568e-03, 2.62187541e-01, 5.55602312e-02, 2.09347799e-01, + 6.81416085e-03, 6.24816194e-02, -1.57044847e-02, 5.38076796e-02, + -1.21786725e-03, 3.13683510e-01, 7.88503736e-02, 1.91961545e-02, + -7.89848214e-04, 2.04859436e-01, 1.77144408e-01, 1.27074376e-01, + -2.01367047e-02, 1.25420019e-01, 6.45905286e-02, 2.66942471e-01, + 3.49235058e-01, 2.52543539e-01, 3.24409813e-01, 2.59198427e-01, + 3.92820060e-01, 1.41132906e-01, 9.48882196e-03, 1.74179897e-01, + 5.39034963e-01, 1.58920869e-01, 2.73100972e-01, 1.81179568e-02, + 1.68145299e-01, 3.43630970e-01, 1.51725076e-02, 2.98200041e-01, + 2.57570058e-01, 2.81464636e-01, 2.47943982e-01, 1.76195994e-01, + 5.23593724e-02, 1.48531556e-01, 2.94840173e-03, 2.51852840e-01, + -7.62507394e-02, 1.64415836e-01, 4.18539464e-01, 1.66102067e-01], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-8.55890103e-03, 3.67313088e-03, -4.36119083e-03, -8.68369080e-03, + -1.34754460e-03, -6.52316585e-03, -5.01648569e-03, -8.97654798e-03, + 3.13179940e-03, -7.83655513e-03, 2.37924187e-03, -5.66111412e-03, + -1.27145583e-02, -8.22409242e-03, 2.74701975e-03, -4.37194854e-03, + -4.04660078e-03, -3.13023920e-03, -1.81288610e-03, -9.03141592e-03, + -1.66141614e-03, -5.17393708e-01, -1.02338977e-02, 3.56379431e-04, + -7.41085410e-03, 9.37872624e-04, -2.64340243e-03, 5.81945898e-03, + -1.19625553e-02, -4.65264637e-03, -6.52810233e-03, -7.84709677e-03, + 1.05619885e-03, -9.87111218e-03, 4.54840425e-04, -5.49622951e-03, + 7.56876124e-03, -5.51949022e-03, -3.83866858e-03, -1.84899215e-02, + 3.52073257e-04, -2.78544938e-03, -5.13466494e-03, -4.08213615e-01, + -5.55376470e-01, -1.75705249e-03, 4.89682052e-03, 1.53476396e-03, + 1.07086054e-03, -3.85978585e-03, 9.40655766e-04, 6.59122784e-03, + -3.99606815e-03, -6.94369990e-03, -2.59699044e-03, 7.18095253e-05, + -8.83713830e-03, -1.96340843e-03, -2.48700310e-03, -9.87609662e-03, + -6.28701295e-04, -7.88388960e-03, 7.09942728e-03, 9.48800612e-03, + -1.95277121e-03, -1.22706648e-02, 2.85065360e-03, 9.00720805e-03, + -7.02000320e-01, 4.25555557e-03, -1.28851901e-03, 2.98675569e-03, + -3.15172016e-03, -1.08107440e-02, -8.73899739e-03, -1.31636495e-02, + -9.91051551e-04, -1.34902084e-02, 3.62220523e-03, 3.06941615e-03, + -1.32400328e-02, -2.95568615e-01, 7.01408414e-03, -5.55180886e-04, + 3.73574020e-03, -7.83216953e-03, 7.60261260e-04, -9.49840934e-04, + -1.34194596e-02, -1.67394914e-02, -8.75537097e-03, -1.25364307e-03, + -6.53209072e-03, -1.00438660e-02, -4.49657347e-03, -1.30805969e-02, + 7.16269715e-03, -1.18865566e-02, 1.15038871e-04, -7.90204573e-03, + -1.13492664e-02, -1.28111634e-02, -2.23936560e-03, -5.98883303e-03, + -1.25457710e-02, -1.07937502e-02, -3.34314862e-03, -7.55358534e-03, + -6.86253095e-03, -2.35536560e-01, -9.40762553e-03, 1.64968718e-03, + -2.74949777e-03, -3.80960264e-05, 1.37860198e-02, 1.36729311e-02, + -1.17542176e-03, -4.75597620e-01, 1.06082670e-02, -4.49829688e-03, + 1.80422806e-03, -3.01245158e-03, -5.77643573e-01, -8.06397665e-03, + 1.37929013e-03, -9.44454595e-03, -2.07981328e-03, 9.04347980e-05, + -5.42783411e-03, 1.13049364e-02, -6.55044988e-03, -9.25507303e-03, + -1.70839131e-02, -3.13228741e-03, -1.04950890e-02, -9.20292269e-03, + -4.24578832e-03, -8.48777406e-03, -1.27782347e-02, -7.17586139e-03, + -5.02940093e-04, -6.85630506e-03, -5.49637014e-03, 5.60814887e-03, + -7.42090074e-03, -1.33384299e-02, -4.22688900e-03, -3.86846159e-03, + -1.31040469e-01, -5.72234035e-01, 1.44373742e-03, -2.45060981e-03, + -5.27091743e-03, -1.61155630e-02, -8.04609526e-03, -3.71574331e-03, + 4.77314461e-03, -2.98696361e-03, -1.80231046e-03, -3.76312784e-03, + 2.70129368e-03, -5.70825338e-01, -1.65502690e-02, -6.61239505e-01, + -9.16241575e-03, -1.41646145e-02, -5.00553148e-03, -4.18055616e-03, + -4.17894637e-03, -1.43527528e-02, 1.20390498e-03, 1.38255325e-03, + -2.05959729e-03, -1.65631052e-03, 8.46141134e-04, -4.89455275e-03, + -3.43874632e-03, -4.29662468e-04, -9.45413113e-03, -4.12517926e-04, + -2.43844930e-03, -8.26839264e-03, -6.06988370e-03, -1.55913727e-02, + -6.00419287e-03, 7.41206226e-04, -4.03943751e-03, -8.32396373e-03, + -1.77192141e-03, -8.27328023e-03, -4.80760727e-03, -1.19567281e-02, + -3.16868187e-04, -1.15141738e-02, 7.74921337e-03, -2.20119534e-03, + 1.03418762e-03, -5.75479446e-03, -4.55250964e-03, -1.04654999e-02, + -6.14567241e-03, -1.54978642e-02, -7.39947148e-03, -1.07835466e-02, + -1.21487593e-02, -6.76209899e-03, 1.25275925e-03, -3.19723099e-01, + -3.41119803e-03, 2.97757820e-03, -5.52376360e-03, -9.89249256e-03, + -4.31803286e-01, -1.36245508e-02, -7.62469275e-03, -1.74430560e-03, + -1.96244684e-03, 5.65344701e-04, -6.50570095e-01, -1.45257404e-03, + -2.76738009e-03, -8.20099376e-03, 2.03965130e-04, -1.13646872e-02, + -2.37802067e-03, -5.39474888e-03, 2.26747943e-04, -2.44784285e-03, + -1.09350355e-02, 8.08892888e-04, -8.32007360e-03, -1.91815558e-03, + -9.69145447e-03, 1.22545648e-03, 4.08804649e-03, -5.84384939e-03, + 3.66018945e-03, 4.44652513e-03, -9.09210462e-03, -4.17519035e-03, + -3.98625523e-01, -3.37073533e-03, -8.35657492e-03, -4.79207002e-03, + -1.39846019e-02, -8.26936215e-03, -8.90700985e-03, -2.52558112e-01, + 3.46473581e-03, -1.71433843e-03, -1.91845442e-03, -3.26462858e-03, + -8.05919431e-03, -5.96231222e-03, -2.60497467e-03, -5.12533821e-03, + -9.94778052e-03, -1.35790939e-02, -7.90751539e-03, -5.66528225e-03, + 2.63936934e-03, -4.14405525e-01, 2.22919116e-04, 8.12191633e-04, + -1.82921547e-04, -7.25358841e-04, -1.08812936e-02, 3.49700474e-03, + -1.17458524e-02, -1.99729693e-03, -5.74411824e-03, -1.04112234e-02, + -5.87427756e-03, 5.42524364e-03, -4.66417568e-03, -3.61002120e-03, + -4.92561430e-01, -5.91339357e-03, -8.21670890e-03, 6.48486370e-04, + -1.16507420e-02, -7.28300912e-03, 9.69227776e-03, -1.40586449e-03, + 1.34252605e-03, -4.55454225e-03, -1.05377687e-02, 8.61465523e-04, + -1.29449442e-02, 3.23367422e-03, -1.80069392e-03, 2.35577201e-04, + -8.91069509e-03, -1.41903665e-03, -8.80321302e-03, -4.76717623e-03, + -2.86917109e-03, -3.21599003e-03, 3.61907203e-03, -7.99222197e-03, + -6.45599188e-03, -5.55047020e-03, -9.93613340e-03, -8.86950456e-03, + 1.25965700e-02, 3.33399116e-03, -1.01383878e-02, 2.34254939e-03, + 5.69945900e-03, -5.75918052e-03, 7.65805482e-04, -1.72046013e-02, + -8.67277035e-04, -1.03675043e-02, -7.15132989e-03, -1.15488358e-02, + -2.48059118e-03, -1.51636070e-02, -1.34787019e-02, -8.78769148e-04, + -6.11743808e-01, -6.02744520e-04, 7.56409660e-04, -3.29326536e-03, + 5.58212679e-03, -9.33979824e-03, -2.45822361e-03, -5.48888650e-03, + -4.64835018e-03, -4.38324921e-03, -3.24640493e-03, -6.00662351e-01, + 1.93342322e-03, -1.04327723e-02, 4.34050598e-04, 3.05471686e-03, + -7.74156442e-03, -1.68287978e-02, -1.33079069e-04, -3.91604519e-03, + -1.49135012e-02, -8.73885211e-03, -5.66676259e-03, 3.75113130e-04, + -5.36525948e-03, -1.39129013e-02, 1.88402925e-03, -6.80929542e-01, + 2.44663190e-03, -1.16351806e-02, -5.38230175e-03, -3.69915483e-03, + -4.62271646e-03, 4.26167302e-04, 1.99855946e-04, -1.86087582e-02, + 7.48794526e-03, -1.61632244e-02, -4.98278052e-01, -8.40996485e-03, + -6.01133797e-03, -8.28853925e-04, -5.63414209e-03, 5.07016899e-03, + -9.90651455e-03, -4.54475218e-03, 1.22021348e-03, 5.08222263e-03, + -5.74093172e-03, 2.58398801e-03, -5.80827938e-03, 2.76929396e-03, + -8.52370542e-03, -8.13993160e-03, -3.10428627e-03, 9.73783247e-03, + -2.26997398e-03, -7.61909829e-03, 4.68078069e-03, -3.60694993e-03, + -8.70567188e-03, -1.44884512e-02, -2.72786571e-03, 4.09194687e-03, + -9.40709934e-03, 4.71121492e-03, -5.04164118e-03, -1.17910588e-02, + 1.36497710e-03, -1.07826991e-02, -1.53774535e-03, -1.12592215e-02, + -2.51332531e-03, -1.03800860e-03, 9.40567255e-03, -1.59978624e-02, + -9.16326977e-03, 5.42400870e-04, -2.22451333e-03, -7.41019659e-03, + -3.90202254e-01, -1.51460301e-02, -6.29314408e-03, -4.74618049e-03, + 3.59613076e-03, -1.17005240e-02, -1.27715422e-02, -1.15625896e-02, + -1.68063808e-02, -1.12238340e-02, -5.49799949e-03, 7.02496106e-03, + -2.54716747e-03, -9.75850690e-03, -3.14950384e-03, -7.65251042e-03, + -1.46619789e-03, 1.40559685e-03, -2.16309223e-02, 4.46918281e-03, + -8.65491666e-03, 5.14270552e-03, 3.51362582e-03, -9.26864985e-03, + -1.17338272e-02, -2.89736595e-03, -1.97741878e-03, -5.52454265e-03, + -8.56413040e-03, -9.34949901e-04, 2.97652930e-03, -2.00355076e-03, + -3.76361725e-03, 1.10566849e-03, -5.63339936e-03, -2.81627057e-03, + -1.68547058e-03, -4.57938435e-03, 6.58686552e-03, -4.55372721e-01, + -5.61221363e-03, -6.51222235e-03, 1.07192027e-05, -8.45061056e-03, + 5.90987783e-03, -5.17501831e-01, -4.53207910e-01, 1.17717020e-03, + -7.01933773e-03, 4.97156382e-03, 1.22358715e-02, -3.75395943e-03, + -1.19871208e-02, 7.76410243e-03, -1.40578160e-02, 3.29062645e-03, + 6.30691764e-04, -1.67252067e-02, 1.95841817e-03, -9.67236888e-03, + -1.17397988e-02, -6.99722324e-04, -6.32376294e-04, -6.49616960e-03, + -2.94766188e-01, 1.18942198e-03, -7.66102644e-03, 3.43931909e-03, + 2.72910809e-03, -5.50179362e-01, -1.15189431e-02, -5.44976583e-03, + -3.17149970e-05, 1.41578354e-03, -6.40626810e-03, -3.47646652e-03, + 6.37558231e-04, -2.99780950e-04, 5.98926889e-03, -4.70549008e-03, + -7.91283231e-03, -9.49787791e-04, -3.55603476e-03, -1.14685595e-02, + -1.00854738e-03, 2.23665382e-03, -2.58625904e-03, -1.55551946e-02, + -9.62937716e-03, -3.64469341e-03, -6.78142011e-01, -3.52118816e-03, + -1.36016486e-02, -1.84007958e-01, -5.86909475e-04, -8.81857052e-03, + -1.06971972e-02, -7.59707938e-04, -5.68926990e-01, -5.22532880e-01, + -2.09414326e-02, -1.72124021e-02, -6.04376057e-03, -3.09983082e-03, + -4.42491425e-03, -8.42403620e-03, 7.68742245e-03, -3.21401260e-03, + 6.66661141e-03, 2.13651289e-03, -4.34801867e-03, -7.12090218e-03], dtype=float32), + kernel: Array([[-0.06747823, 0.12502806, 0.01688226, ..., 0.04402845, + 0.0489017 , 0.02632801], + [-0.01445617, 0.03798477, -0.00418379, ..., 0.01224544, + 0.02048267, 0.01027028], + [-0.01829829, 0.01891246, 0.01610937, ..., 0.03225997, + -0.01639963, -0.03890606], + ..., + [ 0.00887221, -0.00487363, 0.02399558, ..., 0.01346108, + 0.01049829, -0.03240183], + [-0.01674049, -0.00387911, -0.03074075, ..., 0.00146568, + -0.00876305, -0.02836148], + [ 0.01625663, 0.03327096, -0.04843638, ..., 0.08134065, + 0.03410971, 0.01957767]], dtype=float32), + }, + output: { + bias: Array([-0.16341606, 0.2722003 , -0.01811248, 0.01928756, -0.02578784, + -0.10320964, 0.0457683 , 0.00972074, 0.05574008, 0.24268645, + -0.02999464, 0.043408 , -0.0737564 , 0.03339861, -0.04492559, + 0.09947823, 0.05016089, -0.08309133, -0.05127812, 0.11153366, + -0.08190622, -0.06840179, -0.02421227, -0.02895692, -0.04004822, + 0.03941108, -0.0170032 , -0.00097349, 0.07023378, 0.02564083, + 0.03406862, -0.06778747, 0.11675587, -0.08771975, 0.04982556, + -0.07175723, -0.04058652, -0.01989682, -0.02999596, -0.07330333, + 0.03119163, -0.0147863 , -0.83884764, 0.04124971, -0.147959 , + 0.06891607, 0.05300924, 0.2076162 , 0.04646805, -0.05764642, + -0.03448724, 0.05170846, -0.02983713, 0.14581041, -0.04477751, + -0.0028539 , -0.08908532, -0.09743858, 0.05685575, 0.08213867, + 0.0757231 , -0.13601153, -0.02382083, -0.01865633, -0.00175037, + -0.03086867, -0.03948168, 0.06656119, 0.03003574, 0.07914743, + -0.02571616, 0.03044318, 0.05824798, 0.02247159, 0.7948062 , + -0.00407046, -0.07054033, 0.12808433, -0.08461454, 0.02311847, + -0.65620077, 0.10964675, 0.10405939, -0.00352513, -0.02670519, + 0.15050422, 0.13908558, 0.14308207, 0.03576402, -0.11146745, + 0.07615335, 0.04982926, -0.11778215, -0.10201599, -0.01381564, + -0.01693749, -0.03509882, 0.01073222, -0.01447865, 0.08480705, + 0.02910439, 0.00359502, 0.01869445, -0.05797101, 0.09884419, + 0.0161234 , 0.00668302, -0.08439763, -0.13141344, -0.03596077, + -0.14647698, 0.03868678, 0.11478355, -0.14481665, 0.00641343, + 0.16066928, 0.04811645, -0.02036389, -0.03384693, 0.06242058, + -0.00735435, 0.02676932, 0.12175596, -0.06479023, 0.04808705, + 0.05643554, -0.07896087, -0.02733074], dtype=float32), + kernel: Array([[ 0.04431782, -0.07884661, -0.02030264, ..., 0.03209 , + 0.02117256, 0.00271614], + [-0.05527279, 0.03204606, -0.00293706, ..., -0.03798404, + -0.02230972, 0.00913886], + [-0.01904109, -0.02136857, 0.00946268, ..., 0.00186177, + -0.02516458, 0.05516352], + ..., + [-0.03168841, 0.07299381, -0.00359565, ..., -0.01263119, + -0.04168412, 0.01372963], + [-0.02394411, 0.08588913, 0.00119561, ..., 0.0126469 , + -0.04365737, 0.0029881 ], + [-0.01780676, 0.05708533, 0.01049639, ..., 0.03207007, + -0.03454534, 0.0507676 ]], dtype=float32), + }, + }, + }, + }, + 4: { + attention: { + attention: { + key: { + bias: Array([ 0.25172058, 0.68627775, 0.5949247 , 0.1658349 , -0.3173467 , + 0.1774752 , 0.25412992, 0.08010545, 0.56002 , 0.17958656, + -0.17482378, 0.07066374, 0.53656256, -0.5385957 , 0.03753721, + 0.3764516 , -0.05485516, 0.7278333 , -0.7301842 , 0.7662722 , + 0.8557523 , 0.19676152, -0.7459649 , 0.11304809, -0.27465415, + 0.41445032, 0.31016302, -0.9027352 , 0.15300739, -0.12067419, + 0.11333649, 0.36544785, 0.10615282, -0.19266807, 0.00742444, + -0.3133698 , 0.13811715, 0.17462967, -0.04043375, 0.01431966, + 0.02120961, -0.31457424, 0.24998754, -0.30905357, -0.33957723, + 0.44098425, 0.31071022, 0.5162943 , 0.1519571 , 0.20767058, + 0.11702721, 0.30976802, -0.26687843, 0.02610935, -0.38046166, + 0.15602328, 0.1049099 , -0.3816162 , 0.34945863, -0.13488789, + 0.25807512, -0.14170453, -0.2737117 , 0.6794868 , -0.5637404 , + -0.7506761 , -0.59640557, -0.35966128, -0.36779413, 0.6445572 , + -0.15863825, 0.28188404, 0.6817374 , -0.30411932, 0.15068318, + 0.21245846, 0.47254717, -0.3180801 , -0.72213024, -0.00796549, + 0.52852696, -0.74865776, 0.72052324, -0.6231846 , -0.35524404, + -0.6322673 , 0.5927172 , 0.7579426 , 0.22011517, 0.37161058, + -0.58662343, -0.5653432 , 0.4952815 , 0.34009102, 0.4199334 , + -0.04050065, 0.967862 , 0.62645787, -0.24915175, 0.6059676 , + -0.12386217, -0.97235304, 0.98436683, -0.71207863, -0.2554404 , + 0.29857978, -0.6535443 , -0.00581208, 0.54410416, -0.64588374, + -1.2313039 , -0.7612568 , -0.5072008 , -0.15740068, 0.5934992 , + 0.29966694, -0.78861797, 0.95648044, 0.90275365, 0.3240316 , + -0.7827806 , -0.90508664, 0.24878727, -0.92607623, 0.45041922, + -0.7728655 , 0.2784089 , 0.75287944], dtype=float32), + kernel: Array([[ 1.6211687e-01, -2.6181787e-02, -5.1868878e-02, ..., + -2.0115136e-01, -1.6275570e-01, 4.0626042e-02], + [ 2.0827402e-01, -1.1491413e-01, -1.4099246e-01, ..., + 1.2704584e-01, 2.0065814e-01, 3.5636088e-01], + [ 1.4854251e-01, -1.1401314e-01, -4.7746781e-02, ..., + -8.4292240e-02, 1.9939394e-01, 9.9089958e-02], + ..., + [-5.0814073e-03, 7.1649961e-02, 8.1174605e-02, ..., + 1.8439499e-01, -1.8450679e-01, 3.2407336e-04], + [ 5.2984446e-02, -4.0257105e-01, -3.7533784e-01, ..., + 2.7334958e-01, -7.9834215e-02, -2.6604572e-01], + [-1.3808161e-01, -2.5356847e-01, -2.2354984e-01, ..., + -3.9307380e-01, 3.9180964e-01, 4.5249498e-01]], dtype=float32), + }, + query: { + bias: Array([ 3.67379338e-01, -3.86799455e-01, -3.21791142e-01, -8.42057094e-02, + 4.18459654e-01, 5.85961580e-01, -4.60834682e-01, -6.04065180e-01, + 1.94889054e-01, 5.26792884e-01, -4.81182262e-02, -8.55236873e-02, + -5.21454439e-02, 4.12250131e-01, -1.91803545e-01, -5.08220792e-01, + -6.88099623e-01, 1.73875868e-01, -6.42500934e-04, 3.13504860e-02, + 7.57871941e-02, -1.18366331e-01, 7.69717107e-03, 6.50874972e-01, + -1.46498501e-01, -3.85033339e-02, 1.68164968e-01, 2.94075936e-01, + 5.43904722e-01, 1.82771876e-01, 4.28779982e-02, 3.94661844e-01, + 2.40447834e-01, 7.63046503e-01, 1.41106710e-01, 1.50562599e-01, + 7.20623806e-02, -3.31430919e-02, 9.66051295e-02, -4.25018109e-02, + 2.59480290e-02, 5.50236627e-02, 6.59183860e-02, 7.35617638e-01, + 9.20522138e-02, -1.12395249e-01, -2.63706028e-01, -4.15704325e-02, + 9.31539387e-02, -6.19517326e-01, 3.92871583e-03, -2.88900491e-02, + 8.03363383e-01, -1.52911708e-01, 7.79530257e-02, -1.06074944e-01, + -5.33656120e-01, 1.22914307e-01, 1.09775752e-01, 3.54613900e-01, + -4.29570854e-01, 1.18217006e-01, 5.40162444e-01, -1.17120035e-01, + -6.53074026e-01, -4.32521045e-01, -7.02763855e-01, -4.33477342e-01, + -4.24702972e-01, 4.39991415e-01, -2.19614744e-01, 4.09397632e-01, + 5.78722954e-01, 1.95280928e-02, 2.52549767e-01, 5.14830410e-01, + 5.46005607e-01, -6.23379529e-01, -4.24228251e-01, -2.12890625e-01, + 2.95019180e-01, -6.27236724e-01, 5.73115885e-01, -3.63081634e-01, + -3.95359695e-01, -4.50571030e-01, 5.81661999e-01, 5.84261417e-01, + 2.53536344e-01, 5.42065144e-01, -3.84602755e-01, -4.97589141e-01, + 4.21692282e-01, 4.60142285e-01, 3.82059246e-01, -3.72222006e-01, + 5.13080180e-01, 2.00476050e-01, -6.01216495e-01, 8.54436874e-01, + 2.28271738e-01, -3.90953124e-01, -9.38059986e-02, -8.37827742e-01, + 5.72855584e-02, -1.37482524e-01, 1.96730390e-01, 1.36508822e-01, + 1.39595941e-01, -4.52210784e-01, 1.51813626e-01, -7.74807990e-01, + 2.37015381e-01, -3.03821355e-01, -3.49135190e-01, -1.64326906e-01, + -8.14965189e-01, -1.61662132e-01, 8.50251555e-01, 2.36134544e-01, + 1.86818361e-01, -8.37599099e-01, 9.94282067e-02, 8.31496865e-02, + 1.03579015e-01, -4.60613877e-01, 8.01612064e-02, 8.32802534e-01], dtype=float32), + kernel: Array([[ 0.00486563, 0.01157592, 0.0173437 , ..., 0.06944456, + -0.07013498, 0.01022998], + [-0.07428966, -0.05709548, -0.03850973, ..., -0.09689807, + 0.22320649, 0.2770357 ], + [-0.10973258, 0.19551165, 0.12421291, ..., -0.02782791, + -0.36778155, -0.19087008], + ..., + [ 0.30281523, 0.3621789 , 0.32462892, ..., 0.0546729 , + -0.04859495, -0.10934605], + [ 0.5838381 , 0.22978936, 0.17813566, ..., -0.16678724, + 0.70342004, 0.18783332], + [-0.48419687, 0.04412645, -0.01096392, ..., -0.05038855, + 0.37432894, 0.02556431]], dtype=float32), + }, + value: { + bias: Array([ 0.00544928, 0.17689343, 0.04735808, -0.18264541, -0.04237033, + -0.17941093, 0.00340467, -0.00149406, -0.23728922, -0.07449124, + -0.2651422 , -0.0984367 , -0.12018497, -0.08068941, -0.01734875, + -0.06606279, 0.1477764 , -0.02022878, -0.1402983 , -0.02072977, + 0.12049822, 0.06401265, 0.00406305, 0.15095161, -0.20597391, + 0.21405579, -0.08197237, -0.00951357, -0.13306661, -0.08558191, + 0.08653639, -0.10818963, -0.1644184 , -0.15369424, 0.3490528 , + 0.01967676, -0.16458155, 0.10356332, -0.2320539 , -0.0433312 , + 0.294788 , 0.08036903, -0.12358049, 0.16056572, 0.25535643, + 0.16914676, -0.06201077, 0.28501964, -0.13111252, 0.11861207, + 0.26627177, 0.16901444, -0.19237208, -0.08137509, 0.2016543 , + 0.15213989, -0.14299503, -0.24307302, -0.29300013, -0.0722383 , + 0.05268372, -0.14988767, -0.19426815, -0.24243473, -0.06109034, + -0.13096058, 0.155086 , 0.15769799, 0.10678913, 0.00772086, + 0.12339002, -0.02504941, -0.13470462, 0.02844545, 0.09698638, + 0.11937644, -0.01387584, -0.13328613, -0.05046262, -0.02981377, + -0.11246655, -0.06740061, -0.00562667, 0.15140204, 0.08057123, + -0.03175431, 0.05622792, 0.11395832, -0.09259764, -0.02921034, + -0.10965273, 0.09367548, 0.07572406, 0.17622064, -0.12645517, + -0.06927612, -0.13641733, 0.11861874, -0.01573502, -0.10696157, + -0.14588912, -0.16159363, 0.1815869 , 0.08687254, 0.11738776, + -0.0497781 , -0.12963611, 0.01230228, -0.10121961, -0.06638794, + 0.11086646, -0.05197389, 0.24918154, -0.16594617, 0.20172946, + 0.06618202, 0.1124498 , 0.11753738, -0.19368784, 0.04725859, + -0.13396324, -0.10240424, -0.11015446, 0.0015793 , 0.05593257, + -0.12091528, 0.08938907, -0.03860855], dtype=float32), + kernel: Array([[ 0.06369832, 0.04487664, -0.0572863 , ..., -0.00566542, + 0.00958204, -0.0823305 ], + [-0.48305708, -0.05723378, 0.09843583, ..., 0.01919795, + 0.04487755, 0.36410812], + [-0.39578405, 0.17672975, -0.04119537, ..., 0.06564505, + -0.01046039, -0.30802208], + ..., + [-0.04638171, -0.03762753, -0.15475406, ..., 0.11968867, + -0.1331625 , 0.01799979], + [ 0.34840006, 0.0131262 , -0.17438433, ..., 0.21620421, + -0.14971787, -0.73045504], + [ 0.14833876, 0.18964465, 0.05431836, ..., 0.10322906, + 0.10638059, -0.31155896]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-1.46490365e-01, 2.10801542e-01, -1.72206163e-02, -3.77512947e-02, + -7.96984136e-03, -7.52876028e-02, 1.01055026e-01, 2.95142401e-02, + 4.51585725e-02, 2.89253950e-01, -3.78346108e-02, 3.64244655e-02, + -6.14614636e-02, 4.95899320e-02, -3.73293981e-02, 5.79853542e-02, + 8.40439461e-03, -1.02269538e-01, -7.46445209e-02, 1.14303716e-01, + -2.32893787e-02, -3.25526744e-02, -1.71621908e-02, -2.69102436e-02, + -4.24926169e-02, 6.34983480e-02, -9.30675417e-02, -2.39975707e-04, + 8.57859775e-02, 5.52252270e-02, 8.55132714e-02, -9.58903059e-02, + 1.34693608e-01, -1.06799692e-01, 8.42991620e-02, -7.29722008e-02, + -3.16112302e-02, -1.33186892e-01, -3.01891286e-02, -1.23371415e-01, + 2.47585457e-02, 9.32596158e-03, -8.83968234e-01, 3.25804874e-02, + -1.36455193e-01, 5.22013418e-02, 7.06611201e-02, 2.37159967e-01, + 8.91016126e-02, -2.99384054e-02, -1.87199526e-02, 2.26440327e-03, + -2.86240019e-02, 1.45185545e-01, 1.16965972e-01, -3.56947407e-02, + 1.23156169e-02, -7.89485201e-02, 2.42000669e-01, 5.18658943e-02, + 1.23368412e-01, -1.33664429e-01, -5.28409593e-02, -1.49152186e-02, + -7.04273768e-03, 4.40860493e-03, -9.83215496e-02, 4.80295867e-02, + 1.94351804e-02, 9.06426013e-02, -1.53706968e-02, 5.25671393e-02, + 7.91645125e-02, 2.07034983e-02, 7.01766312e-01, 1.80371366e-02, + -1.23936214e-01, 3.77195626e-02, -9.12794545e-02, -5.07131638e-03, + -5.29812515e-01, 2.14143340e-02, 7.39834160e-02, -7.41400057e-03, + 6.01231493e-02, 7.33223036e-02, 1.53090119e-01, 1.70849040e-01, + -4.07319590e-02, -4.82572988e-03, 2.90474202e-02, 4.86663431e-02, + -1.04625016e-01, -8.78558010e-02, -4.20586728e-02, 3.25415730e-02, + -7.98427388e-02, -1.87286697e-02, -2.41589937e-02, 8.27142298e-02, + 4.98151742e-02, 5.18750064e-02, 3.29844393e-02, -3.68855074e-02, + 8.92899334e-02, -4.56710160e-03, 2.97619849e-02, -7.04986304e-02, + -1.08520284e-01, -1.11746334e-01, -7.71660432e-02, 8.49972386e-03, + 2.37586319e-01, -1.24874555e-01, -1.93465855e-02, 8.61749053e-02, + 8.92835855e-02, -8.50701034e-02, -3.94290835e-02, -2.06257198e-02, + -1.92599650e-02, 2.76334919e-02, 6.91349506e-02, -2.97043975e-02, + 8.56228359e-03, 5.08836359e-02, -1.17148422e-02, -2.16257907e-02], dtype=float32), + kernel: Array([[ 0.16482222, -0.22908561, 0.5409443 , ..., 0.52288747, + 0.45133135, -0.3148359 ], + [ 0.04791335, -0.19366321, 0.02344365, ..., 0.25754645, + 0.15889917, -0.26091245], + [ 0.21047738, 0.00189075, -0.053845 , ..., -0.10111952, + -0.42335704, -0.05204519], + ..., + [-0.07512332, 0.07116441, -0.20204972, ..., -0.3296345 , + 0.2802311 , -0.00193282], + [-0.01015196, 0.04124757, 0.08690923, ..., 0.38333163, + -0.18434097, -0.01834725], + [ 0.0501605 , 0.7881087 , 0.21703115, ..., -1.006583 , + -0.5440753 , 0.06839034]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([ 7.00707082e-03, -2.13932544e-01, -4.38091122e-02, -1.80086847e-02, + 1.56905562e-01, 3.25518586e-02, -6.51721493e-04, -3.70227383e-04, + -9.39783454e-03, -1.21371552e-01, 6.61288295e-03, -3.71838771e-02, + 1.69692468e-02, -1.13847349e-02, -1.44199710e-02, -1.44837052e-03, + 1.02954609e-02, 7.14356266e-03, -2.28633024e-02, -9.65952054e-02, + -7.08255393e-05, 4.81518777e-03, -1.34146630e-05, 3.52103100e-03, + -1.46182910e-01, -2.66583692e-02, -2.81199813e-02, 2.20761336e-02, + 8.75673592e-02, 1.68776400e-02, 7.70330988e-03, -7.63349526e-04, + 1.69382174e-03, -5.38112130e-03, -2.32931431e-02, 1.33482378e-03, + 5.63706718e-02, -1.53671801e-01, -5.60674816e-04, 1.17010390e-02, + 4.67808433e-02, -1.00002938e-03, 7.54884481e-02, 9.78838280e-03, + -2.51426380e-02, -1.93388745e-01, -1.30395563e-02, -3.22579712e-01, + -1.23171741e-02, -2.32061837e-02, -5.90802915e-02, 6.99069537e-03, + 1.86737273e-02, -1.55908689e-01, 5.42053096e-02, -1.68696530e-02, + -3.76862218e-03, 1.75712314e-02, -2.17723008e-02, -1.39841557e-01, + -1.11271013e-02, -8.20189528e-03, -1.91661697e-02, -2.04071011e-02, + 5.92404092e-03, 2.63817403e-02, 2.42956262e-02, 6.90133497e-02, + -7.82949443e-04, 6.91150129e-03, -6.76295608e-02, -8.51555727e-03, + -1.79697247e-03, 4.60966006e-02, -4.87765044e-01, 3.24066654e-02, + -7.08101019e-02, -2.59833187e-02, -4.40584077e-03, 3.93147208e-03, + 2.80292213e-01, -3.62146534e-02, 1.14013357e-02, 1.44051407e-02, + 8.79329368e-02, -3.40585620e-03, 1.61548674e-01, -4.32909727e-02, + 5.22663780e-02, -8.88153631e-03, -5.72965667e-03, 1.60463457e-03, + -1.44764967e-03, -4.41680253e-02, -6.15086630e-02, 2.85431091e-03, + -5.11140227e-02, 3.58232483e-02, 3.14552360e-03, -8.34190473e-02, + -6.30146563e-02, 3.52296643e-02, 1.75646115e-02, -2.81428844e-01, + 1.00661494e-01, -1.63044967e-02, -1.77673250e-03, -2.50048609e-03, + 9.31597501e-02, -5.21637946e-02, -1.21001536e-02, -3.70397139e-03, + -1.89869609e-02, 4.85324338e-02, -1.78271660e-03, -1.17816813e-02, + 6.07670583e-02, -6.83511868e-02, -7.59183466e-02, 2.76686642e-02, + -2.92036515e-02, 7.23615754e-03, -1.44946417e-02, 6.89900964e-02, + -1.19854533e-03, -3.26084495e-02, -3.72714689e-03, 8.10674950e-02], dtype=float32), + scale: Array([ 0.01037673, 0.20990598, 0.17576548, 0.21721113, 0.29847282, + 0.69343257, 0.00376189, -0.01851019, 0.03572948, 0.19570692, + -0.02536977, 0.42987674, 0.01018215, 0.15321532, 0.48395535, + -0.00318469, 0.16068086, 0.01112385, 0.36739224, 0.4355385 , + -0.01194211, -0.00359833, -0.086844 , 0.00597237, 0.5273556 , + 0.03060911, 0.6654633 , -0.05266754, 0.36045882, 0.10801558, + 0.29989985, -0.04426623, 0.00284036, 0.02850281, 0.08200129, + 0.00828265, 0.4786764 , 0.32044822, 0.6007385 , 0.39719838, + 0.38223657, -0.02341203, 0.02303379, 0.25033328, -0.08627673, + 0.5241369 , 0.544966 , 0.3724897 , -0.32547915, -0.24456456, + 0.32741693, 0.11249542, 0.07124001, 0.5731743 , 0.22054233, + 0.68054676, -0.04105046, 0.07573894, 0.05682403, 0.4511571 , + 0.51980895, 0.05435016, 0.06856246, 0.16396722, 0.00960637, + 0.06122901, 0.23258007, 0.3978171 , 0.23067153, -0.0107969 , + 0.517839 , 0.03140327, 0.02615155, 0.70969087, 0.1419986 , + 0.37057772, 0.54277253, 0.18082201, 0.0020515 , -0.11159092, + 0.43598834, 0.37816525, 0.586952 , 0.35433614, 0.50891733, + 0.00721514, 0.27681717, 0.38934416, 0.4676223 , 0.18506213, + 0.18672895, 0.01464603, -0.00548072, 0.72455907, 0.26698393, + -0.00541573, 0.8330279 , 0.64272374, 0.00286223, 0.52491736, + 0.5533957 , 0.3439355 , 0.32210883, 0.56023586, 0.5354366 , + 0.3412869 , 0.05528033, -0.15407614, 0.28940812, 0.39378768, + 0.5161475 , -0.03409952, 0.05442618, 0.390517 , 0.00593842, + 0.06129869, 0.42658743, 0.4008259 , 0.46517026, 0.30656123, + 0.2991743 , 0.03482211, 0.5670987 , 0.35471496, 0.08178211, + 0.0831662 , 0.00438205, 0.29678687], dtype=float32), + }, + layernorm_before: { + bias: Array([-1.00732700e-03, -3.65120143e-01, -3.34881805e-02, 4.26957064e-04, + -7.94777123e-04, -6.86906278e-03, -1.02201905e-02, -4.80189286e-02, + 1.55887365e-01, -1.59472153e-01, 5.88005059e-04, -4.74225767e-02, + 1.99526129e-03, -6.20241184e-03, 2.79833134e-02, 7.60109648e-02, + -4.03508246e-02, -5.76728676e-03, -8.25315190e-04, -5.45392111e-02, + -9.78699885e-03, -2.21987873e-01, 8.89718998e-03, 1.56212691e-02, + 2.51562544e-03, -2.31638588e-02, 2.57532690e-02, -3.37633025e-03, + -4.75819921e-03, 2.06635101e-03, 4.15294990e-02, -1.22549701e-02, + 9.92366020e-03, 1.33297034e-02, -6.88615516e-02, -2.13969056e-03, + -9.33400914e-03, -1.47152513e-01, 1.47844050e-02, 2.68603060e-02, + 1.06496513e-01, -1.40187237e-03, 2.46544197e-01, -1.18307155e-02, + 1.22905318e-02, 7.02085486e-03, 3.16104256e-02, 1.29607348e-02, + -1.19067011e-02, 5.89615945e-03, 1.45399803e-02, 7.92596638e-02, + 8.44230223e-03, -5.16842827e-02, 2.47099735e-02, 7.68707599e-03, + -7.68995145e-03, 3.25880721e-02, -3.79534140e-02, -1.32843507e-02, + 1.00192919e-01, 8.90694335e-02, 1.43536124e-02, -1.85336508e-02, + 4.20249486e-03, 4.04581055e-02, 1.56428572e-02, 5.54322731e-03, + 7.59444525e-03, -9.50715330e-04, -5.35057345e-03, 2.70144927e-04, + 1.31974250e-01, 2.08387990e-02, -5.05078256e-01, -4.19664755e-02, + 2.06986954e-03, 2.65133400e-02, 1.08654350e-02, -2.48142332e-02, + -2.32796840e-04, -3.82554834e-03, -3.06907180e-03, -3.17831822e-02, + -2.74789296e-02, 1.76229049e-04, -5.93862869e-03, -2.54624914e-02, + 5.30235134e-02, 3.77030671e-02, -3.30569521e-02, 6.81126863e-02, + -3.34922522e-02, -5.41602597e-02, -4.98229451e-03, -1.29764492e-03, + 1.27451606e-02, 5.30401291e-03, 8.72196928e-02, 5.26771834e-03, + -5.91456294e-02, 2.01584827e-02, -1.72205925e-01, -2.79393774e-02, + -3.19484659e-02, -3.31652700e-03, 3.22748870e-02, 1.10277152e-02, + 7.98177347e-02, -2.52419740e-01, 6.41727597e-02, 2.78249923e-02, + -1.18425563e-01, 5.46144024e-02, -1.38351889e-02, 1.97037379e-03, + 3.38066697e-01, -1.05525143e-02, -5.22397608e-02, -8.39580037e-03, + 1.55923087e-02, 2.35547218e-02, -4.73677404e-02, 7.04186484e-02, + -5.95991965e-03, 5.85895148e-04, 4.35512178e-02, -9.74338278e-02], dtype=float32), + scale: Array([-0.00892128, 0.22834587, 0.07191961, -0.03907548, -0.00996873, + 0.35814652, 0.27738327, 0.11775097, 0.07957964, 0.2968416 , + 0.01412186, 0.27139878, -0.00244969, -0.01754642, 0.31717873, + 0.24255775, 0.34514305, -0.0094078 , -0.01727619, 0.2070495 , + 0.23165925, 0.48358163, 0.08980764, 0.16053861, 0.07536255, + 0.02379382, 0.3457654 , 0.22846563, -0.1335053 , 0.01941462, + 0.26780826, 0.09502309, 0.14652625, 0.05376663, 0.22817616, + 0.20199119, 0.2713218 , 0.26872477, -0.0762281 , 0.22205614, + 0.31939635, 0.24618064, 0.17259929, 0.21590437, 0.02615085, + 0.17664434, 0.36018652, 0.12768349, 0.20409435, 0.14295201, + 0.18369466, 0.3224867 , 0.01613115, 0.11641789, 0.19955999, + 0.08757862, -0.01346728, 0.198473 , 0.29986396, 0.36929435, + 0.5931218 , 0.19321899, 0.11543731, 0.16761477, -0.00143133, + 0.07413565, 0.2747692 , 0.02911899, 0.12630646, 0.02624349, + 0.3049502 , 0.17227921, 0.35921425, 0.295701 , 0.13059482, + 0.19860806, 0.01733009, 0.24637464, 0.00986639, 0.19075045, + -0.00124613, 0.27278668, 0.01111688, 0.16424142, 0.18801676, + 0.01323086, 0.00368139, 0.03970122, 0.26844364, 0.10092615, + 0.09264332, 0.16377948, 0.22653754, 0.35064098, -0.05561429, + -0.00500794, 0.07151927, 0.0050664 , 0.2468497 , 0.10480628, + 0.17903912, 0.08759426, 0.41523352, 0.24769782, 0.11629754, + 0.03568345, 0.20900775, 0.05634574, 0.4340352 , 0.3333243 , + 0.41587278, 0.20580357, 0.3584831 , 0.2544934 , 0.01131326, + 0.34364027, 0.22543208, 0.13177043, 0.19698657, 0.21796568, + -0.03295131, 0.298913 , 0.47042173, 0.26397368, 0.02212258, + 0.01400256, 0.3371711 , 0.22812258], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-5.27199693e-02, -1.22556537e-02, -4.80708107e-02, -2.20634907e-01, + -1.40253054e-02, -4.87366840e-02, -8.70805699e-03, -1.54133514e-02, + -2.52253171e-02, -3.50418165e-02, -2.73905657e-02, -3.95800360e-02, + -3.71573754e-02, -2.81724222e-02, -1.33444788e-02, -1.28413305e-01, + -4.53637205e-02, -3.11292056e-02, -3.32505777e-02, -3.26416194e-02, + -3.45684476e-02, -1.40969567e-02, -2.91586872e-02, -2.76421569e-02, + -3.73606943e-02, -3.91720347e-02, -3.77491489e-02, -1.35208545e-02, + -2.09600553e-02, -4.55045812e-02, -5.19248992e-02, -2.19188482e-02, + -1.82508994e-02, -6.17397487e-01, -3.09374705e-02, -5.18258326e-02, + -2.61735450e-02, -3.89497764e-02, -1.46954479e-02, -1.18027776e-02, + -2.21383367e-02, -3.07989083e-02, -6.69780612e-01, -1.84117928e-02, + -4.12670113e-02, -8.98960326e-03, -3.24395895e-02, -1.61407497e-02, + -4.12646681e-02, -5.37053635e-03, -3.38159055e-02, -3.29089276e-02, + -1.68235525e-02, -3.26590054e-02, -2.52911523e-02, -4.61294129e-02, + -5.67609370e-01, -2.84224618e-02, -4.13801000e-02, -3.34288217e-02, + -3.53450403e-02, -3.29895802e-02, -4.02423888e-02, -2.35914946e-01, + -4.60832685e-01, -3.70765813e-02, -2.73666903e-02, -2.54514460e-02, + -1.44379437e-02, -4.47799824e-02, -4.39045168e-02, -4.60769348e-02, + -1.07184583e-02, -2.94511896e-02, -3.48258428e-02, -4.94903252e-02, + -4.71179605e-01, -3.04534864e-02, -2.14969758e-02, -2.71185916e-02, + -2.13122219e-02, -2.35180482e-02, -3.83645929e-02, -2.21622270e-02, + -2.68899292e-01, -3.47404443e-02, -3.07555981e-02, -2.66567357e-02, + -2.49998886e-02, -4.09555882e-01, -3.07542346e-02, -1.73584837e-02, + -3.32305208e-02, -3.31670046e-02, -3.28688100e-02, -2.02093106e-02, + -3.04487552e-02, -4.77275476e-02, -4.22687948e-01, -4.02031802e-02, + -2.18092110e-02, -3.58968377e-02, -2.14504618e-02, -3.14717107e-02, + -6.16603434e-01, -3.43027376e-02, -4.16996516e-02, -2.88160015e-02, + -3.17607224e-02, -4.91070986e-01, -4.11182418e-02, -3.82331535e-02, + -1.99981369e-02, -3.31944197e-01, -1.49236154e-02, -4.61634010e-01, + -4.73840237e-02, -1.24957161e-02, -3.81090641e-01, -1.33936629e-01, + -4.61227804e-01, -3.94086614e-02, -5.02103120e-02, -1.19854463e-02, + -2.26863697e-02, -3.57204154e-02, -2.89868433e-02, -3.16134878e-02, + -2.95881126e-02, -3.33401635e-02, -2.84984969e-02, -3.39026898e-02, + -4.84661758e-01, -3.00496016e-02, -2.10581459e-02, -3.03259753e-02, + -2.73086485e-02, -2.92658973e-02, -3.28680202e-02, -4.29259539e-02, + -1.15742823e-02, -3.06628253e-02, -3.49430703e-02, -4.13970686e-02, + -4.64555889e-01, -2.44929623e-02, -2.32402366e-02, -4.51661311e-02, + -3.42191793e-02, -1.04636485e-02, -2.76410598e-02, -2.72743534e-02, + -3.71849462e-02, -2.61768941e-02, -2.64907163e-02, -4.13466580e-02, + -2.31192280e-02, -2.66561229e-02, -2.44059507e-02, -2.19917297e-02, + -3.10962517e-02, -3.92389782e-02, -2.89020780e-02, -4.72296998e-02, + -1.47718657e-02, -4.71781231e-02, -3.74658257e-02, -3.74986231e-02, + -3.42181958e-02, -3.34076025e-02, -1.90843388e-01, -3.28308083e-02, + -1.35448873e-02, -3.84903140e-02, -3.14540192e-02, 9.85380495e-04, + -2.98315007e-02, -4.16581221e-02, -2.44194381e-02, -2.18599793e-02, + -4.54461351e-02, -4.33417000e-02, -2.37555113e-02, -2.65070703e-02, + -3.74873281e-02, -2.00963188e-02, -5.27748942e-01, -1.71016864e-02, + -2.98697986e-02, -3.35949808e-02, -4.21402045e-02, -4.34817523e-02, + -3.61697711e-02, -3.83502319e-02, -2.63728369e-02, -9.91544034e-03, + -4.13909107e-02, -3.27453129e-02, -1.61904916e-02, -2.92313341e-02, + -1.82820782e-02, -3.27837802e-02, -4.96267676e-02, -2.13364977e-02, + -3.07123903e-02, -1.29112750e-02, -3.05922143e-02, -1.34589160e-02, + -1.97699457e-01, -1.26856491e-02, -3.33952643e-02, 3.23457294e-03, + -2.03445368e-02, -2.09414233e-02, -1.26234703e-02, -3.02408133e-02, + -2.34196670e-02, -1.23696076e-02, -2.40369998e-02, -2.81809513e-02, + -2.13961061e-02, -5.47036016e-03, -4.53693084e-02, -3.31229195e-02, + -1.50107536e-02, -2.91419178e-02, -2.07277276e-02, -5.49518093e-02, + -2.95787137e-02, -2.38886494e-02, -2.65640840e-02, -2.46801712e-02, + -3.84779572e-02, -1.79381035e-02, -2.02893633e-02, -3.30808014e-02, + -9.46734101e-03, -3.95300575e-02, -3.69772725e-02, -3.11741419e-02, + -4.37073298e-02, -2.90284753e-02, -3.63519609e-01, -1.85389761e-02, + -1.82310361e-02, -1.51738618e-02, -1.89224500e-02, -2.03151274e-02, + -3.81886885e-02, -2.58216038e-02, -3.81475925e-01, -2.79287975e-02, + -1.52051644e-02, -2.38528401e-02, -4.04975682e-01, -3.95358130e-02, + -4.14295375e-01, -4.38631624e-02, -4.29423749e-02, -3.17317769e-02, + -3.27443145e-02, -2.63639372e-02, -3.31052691e-02, 4.43929108e-03, + -3.27373035e-02, -2.39749160e-02, -3.15703414e-02, -5.37155151e-01, + -8.51269998e-03, -3.57643776e-02, -3.98143642e-02, -3.45447585e-02, + -2.49514170e-02, -2.87994146e-01, -5.97163796e-01, -3.36527936e-02, + -3.88488770e-02, -3.22511718e-02, -4.89461757e-02, -3.97258699e-02, + -2.26812717e-02, -7.74491072e-01, -5.23366332e-01, -4.26757559e-02, + -3.48140113e-02, -2.94348393e-02, -3.54919247e-02, -2.96318009e-02, + -1.08983302e-02, -3.76659594e-02, -1.47500113e-02, -2.93985605e-02, + -2.53436007e-02, 5.79031697e-03, -2.69875340e-02, -2.84397658e-02, + -3.33693698e-02, -4.03800532e-02, -2.46412121e-02, -2.71790083e-02, + -6.01157732e-02, -3.36032584e-02, -3.37207466e-02, -1.18618924e-02, + -4.47471261e-01, -4.70100343e-02, -1.94270499e-02, -1.41834235e-02, + -4.12988141e-02, -4.31584977e-02, -2.58529224e-02, -3.11732627e-02, + -3.12253926e-02, -4.00089920e-02, -2.82256324e-02, -1.63722355e-02, + -2.23478712e-02, -3.37487720e-02, -1.17873568e-02, -3.71523947e-02, + -3.02173067e-02, -2.47102734e-02, -1.75078567e-02, -8.67900550e-01, + -4.19426337e-02, -2.89364811e-02, -3.62409726e-02, -1.40926614e-02, + -3.56511444e-01, -3.45515162e-02, -5.06763607e-02, -2.30538756e-01, + -3.77509855e-02, -3.81864123e-02, -3.19356658e-02, -2.84604561e-02, + -4.72619906e-02, -1.65800638e-02, -4.46446165e-02, -1.26604997e-02, + -3.79219428e-02, -3.04708660e-01, -3.42885070e-02, -4.41885963e-02, + -3.81699838e-02, -3.03119402e-02, -2.27376912e-02, -1.89236663e-02, + -4.71340455e-02, -1.56659111e-02, -3.82196195e-02, -3.61527875e-02, + -2.05292553e-02, -3.56830619e-02, 3.07818269e-03, -7.64196873e-01, + -2.73544770e-02, -5.15536249e-01, -1.76008549e-02, -6.18065059e-01, + -3.16055007e-02, -4.54795063e-02, -2.52228677e-02, -1.78863723e-02, + -3.71731035e-02, -3.15613933e-02, -3.33131254e-02, -1.16747636e-02, + -1.52025633e-02, -3.18961823e-03, -2.85964645e-02, -3.42308469e-02, + -6.01763930e-03, -1.44698163e-02, -3.33204381e-02, -1.46625296e-03, + -1.37048895e-02, -2.66424436e-02, -4.91908789e-02, -3.08587961e-02, + -2.29315758e-02, -3.17254931e-01, -1.58655606e-02, -8.84249713e-03, + -1.90528110e-02, -4.16713282e-02, -3.05165984e-02, -3.23779024e-02, + 1.71483532e-02, -1.71599649e-02, -1.27734728e-02, -3.09068114e-02, + -3.11801564e-02, -3.09810303e-02, -3.29174399e-02, -4.35743593e-02, + -3.38067338e-02, -2.26877723e-02, -3.14960629e-02, -2.18670722e-02, + -2.10534912e-02, -2.37145834e-02, -3.18468064e-02, -3.17289941e-02, + -3.52053195e-02, -3.62415873e-02, -3.78505550e-02, -8.54883492e-01, + -6.25598431e-01, -2.91614741e-01, -4.12686355e-02, -3.40597779e-02, + -4.49921228e-02, -1.83722675e-01, -6.20228611e-03, -9.09710955e-03, + -4.55179624e-02, -4.24469262e-02, -4.34849598e-02, -1.18993660e-02, + -2.96284892e-02, -1.82864573e-02, -1.58976540e-02, -4.17148732e-02, + -1.04806479e-02, -2.51498409e-02, -3.05920038e-02, -2.39100829e-02, + -2.05667838e-02, -4.11339812e-02, -2.58969050e-02, -2.54813936e-02, + -3.46504189e-02, -3.05222347e-02, -3.66321839e-02, -1.69235300e-02, + -3.10963802e-02, -2.00775322e-02, -3.57228443e-02, -4.01586331e-02, + -2.64529884e-02, -4.16483805e-02, -4.06530350e-02, -2.00617034e-02, + -4.20974828e-02, -3.35381776e-02, -4.97709587e-02, -1.51588218e-02, + -4.78560105e-02, -1.61089674e-02, -2.51727998e-02, -4.55499999e-02, + -3.32853608e-02, -9.99737252e-03, -9.59883165e-03, -3.51412706e-02, + -3.77242006e-02, -1.95016433e-02, -8.43567073e-01, -1.19444542e-02, + -1.49743442e-04, -3.27145346e-02, -3.28630805e-02, -1.98399983e-02, + -1.64248962e-02, -3.85820344e-02, -1.64614365e-01, -1.02842608e-02, + -1.50990197e-02, -3.22756879e-02, -1.80973206e-02, -2.27656811e-02, + -4.72140424e-02, -3.06061078e-02, -5.52861243e-02, -3.98167334e-02, + -3.09152249e-02, -3.28242555e-02, -3.62482332e-02, -2.84483526e-02, + -7.59550452e-01, -4.58953902e-02, -2.13591140e-02, -2.78457664e-02, + -1.19789401e-02, -2.98056193e-02, -3.15071791e-02, 5.63531043e-03, + -2.39895564e-02, -3.00265819e-01, -4.11561355e-02, -2.46453267e-02, + -4.19642702e-02, -5.44620812e-01, -3.15308571e-02, -3.95948440e-02, + -2.19081920e-02, -2.35356446e-02, -1.36993527e-02, -1.34584429e-02, + -5.07402495e-02, -3.16087119e-02, -2.87754349e-02, -5.01885235e-01, + -4.85699102e-02, -4.50458750e-02, -2.57593673e-02, -2.32719798e-02, + -3.24436370e-03, -2.84067634e-02, -2.62353010e-02, -3.73448022e-02], dtype=float32), + kernel: Array([[ 0.03366351, -0.01376447, 0.00334523, ..., -0.01204627, + 0.0351628 , 0.0058507 ], + [ 0.01479752, -0.00059022, 0.02888848, ..., 0.01372632, + 0.05084606, -0.00852082], + [-0.08748913, 0.06752201, -0.02948284, ..., -0.04957195, + -0.01602156, -0.0564839 ], + ..., + [ 0.03158367, -0.0024418 , -0.01248567, ..., 0.00120247, + -0.03067757, 0.05438066], + [ 0.00710372, -0.07475934, -0.03510208, ..., -0.00362264, + -0.01655151, 0.0695975 ], + [ 0.03609428, -0.04006155, -0.0374678 , ..., -0.02961577, + -0.03186451, 0.05374904]], dtype=float32), + }, + output: { + bias: Array([-1.57430485e-01, 2.07714707e-01, -1.93849746e-02, -2.16451492e-02, + -2.94376239e-02, -1.40355468e-01, 1.01594821e-01, 9.77267139e-03, + 3.38900611e-02, 3.08969438e-01, -4.50306423e-02, 4.16773744e-02, + -5.64103350e-02, 4.65862826e-02, -1.98888164e-02, 5.25933765e-02, + 5.03250724e-03, -9.93339792e-02, -9.45335254e-02, 1.34056509e-01, + -2.08574124e-02, -2.76607983e-02, -2.45273374e-02, -2.30302606e-02, + 4.72312383e-02, 6.59246966e-02, -6.18701801e-02, -2.54601543e-03, + 8.12924057e-02, 4.04679216e-02, 7.19875395e-02, -8.80277827e-02, + 1.21765286e-01, -1.06022269e-01, 8.69699791e-02, -7.91017562e-02, + -7.66793415e-02, -9.08406228e-02, -1.50119308e-02, -1.08254775e-01, + 1.43383909e-02, 1.25184851e-02, -8.86707306e-01, 3.01385261e-02, + -1.39143646e-01, 1.25063658e-01, 5.04919551e-02, 2.49843255e-01, + 5.21366931e-02, -3.97394188e-02, -1.14855161e-02, -1.06901973e-02, + -2.97776181e-02, 1.87453896e-01, 8.37227628e-02, -1.05806980e-02, + -2.71216761e-02, -8.99265856e-02, 2.36086398e-01, 1.00231051e-01, + 1.18218437e-01, -1.38503894e-01, -4.76376042e-02, -2.26873197e-02, + -9.81818512e-03, -6.29501976e-03, -1.04944162e-01, -4.35994677e-02, + 1.57965776e-02, 8.73260945e-02, 8.73774476e-03, 5.19699082e-02, + 7.45332241e-02, -8.56231526e-03, 7.00146556e-01, 9.85437538e-03, + -9.72239524e-02, 3.94646786e-02, -9.50414538e-02, -3.36449943e-04, + -4.73328263e-01, 4.33811136e-02, 1.06056370e-01, -4.22528051e-02, + 2.30107084e-02, 1.08815312e-01, 1.21224724e-01, 1.74949020e-01, + -2.51618456e-02, -2.49235518e-02, 2.24952735e-02, 4.77656983e-02, + -1.08213760e-01, -1.28404409e-01, -3.41621116e-02, 3.01731993e-02, + -7.27362260e-02, -1.10188257e-02, -2.82108001e-02, 1.34670347e-01, + 5.28050587e-02, 5.32444753e-02, 3.20014991e-02, 5.04108816e-02, + -3.19805578e-04, -1.66307837e-02, 3.41970064e-02, -6.60975128e-02, + -1.23339623e-01, -9.24760997e-02, -8.67600888e-02, 8.86741187e-03, + 2.36912012e-01, -1.58255234e-01, -1.41610643e-02, 8.07640478e-02, + 2.81964112e-02, -8.16835240e-02, -1.64333303e-02, -2.98798494e-02, + -7.30156377e-02, 2.87766512e-02, 1.21695735e-01, -5.90394363e-02, + 2.60551348e-02, 5.38477339e-02, -2.07343251e-02, -6.99748248e-02], dtype=float32), + kernel: Array([[ 0.02169053, 0.02691332, -0.00372989, ..., -0.00424074, + 0.03661725, 0.08953198], + [-0.01738822, 0.01985764, 0.03393567, ..., -0.00211225, + 0.01224655, 0.03731351], + [ 0.00689401, 0.02467112, 0.01469113, ..., -0.01584834, + -0.03345798, 0.01720266], + ..., + [ 0.01760239, 0.00760174, 0.01594303, ..., 0.00213138, + 0.00272581, -0.03686259], + [ 0.01555281, -0.08100978, 0.04178749, ..., -0.0257984 , + -0.04736535, -0.03686064], + [-0.01932741, 0.05431129, -0.04024049, ..., -0.00516084, + -0.01021444, 0.02485615]], dtype=float32), + }, + }, + }, + }, + 5: { + attention: { + attention: { + key: { + bias: Array([ 0.15380459, -0.23325492, 0.5298237 , -0.25394794, -0.1792779 , + 0.2635233 , -0.2518766 , 0.51324993, 0.44011664, -0.15511836, + 0.1065289 , -0.18103923, -0.24789087, 0.01313723, -0.00924131, + 0.2846446 , -0.02708824, 0.01480816, 0.12186062, -0.18671922, + 0.47530514, -0.44773427, -0.42629737, -0.25819948, 0.15633082, + -0.48324105, -0.2818154 , -0.61912835, -0.29742548, 0.28758916, + 0.07666907, 0.27997455, -0.5159455 , -0.00791557, 0.40676025, + -0.13460469, 0.5236954 , -0.5017536 , 0.5868691 , 0.01062866, + -0.10968728, -0.03771151, -0.33758804, -0.53288084, -0.6708519 , + 0.60964614, 0.14677066, -0.52277976, -0.00805334, -0.3969676 , + 0.5209981 , 0.34777972, 0.5903865 , -0.7917442 , -0.64988923, + -0.63839513, -0.74101704, 0.11602865, 0.8335415 , 0.2829763 , + -0.41182905, -0.3442916 , -0.47213373, 0.5067961 , 0.05412646, + 0.11681292, 0.26332203, -0.25662488, 0.4626192 , -0.2888915 , + 0.3559262 , 0.43773538, -0.4578352 , 0.32878822, -0.2354647 , + -0.39989737, -0.21824938, 0.23775657, 0.28870043, -0.35698473, + 0.2748483 , 0.05920567, -0.09002119, -0.35282713, 0.20067027, + -0.54461175, -0.23671876, -0.26704383, -0.06359763, 0.01624733, + -0.20986806, -0.0469152 , -0.21808156, 0.10429718, 0.28616387, + -0.0075318 , 0.41809165, 0.21852025, -0.68503106, 0.15126543, + -0.5385752 , -0.6387902 , -0.09934053, 1.1725092 , 0.9802906 , + -0.24173768, -0.3697774 , 0.30808112, 1.2169337 , -0.4494893 , + 0.80163956, 0.62732154, -0.03264387, 0.67577094, -0.39557546, + -0.40416336, 0.0219003 , 0.2829875 , -0.28239524, 0.14352222, + -0.8248839 , 0.02846357, 0.13515349, -0.19850218, -0.22375287, + 0.2135571 , -0.27206707, 0.76276755], dtype=float32), + kernel: Array([[ 0.26200238, 0.13167961, -0.11595184, ..., 0.0414915 , + 0.11238165, -0.30494452], + [ 0.22360708, 0.08118725, 0.02069189, ..., 0.00547806, + -0.10924851, 0.26473263], + [-0.16265303, 0.17587142, 0.12479865, ..., -0.07595218, + 0.1396125 , -0.19043295], + ..., + [-0.26119086, -0.38923773, 0.35491627, ..., -0.7368683 , + -0.49823397, -0.6397334 ], + [ 0.13432133, -0.16010363, 0.22101043, ..., -0.06990113, + 0.25519368, 0.09977636], + [-0.12494502, -0.38168752, 0.16611753, ..., 0.36234066, + -0.0669388 , 0.3607124 ]], dtype=float32), + }, + query: { + bias: Array([-0.18479703, -0.35146004, 0.7537319 , -0.17177609, 0.01369478, + 0.18720719, -0.17976424, 0.15347035, 0.08994126, 0.22886601, + 0.18279251, -0.35825828, -0.06468849, -0.12611765, -0.028998 , + 0.9971776 , -0.18668026, -0.04904906, -0.14856422, -0.24965602, + 0.05141901, -1.1571982 , 0.05819758, 0.14917302, 0.15365446, + -0.85294926, -0.08199202, -0.9097335 , 0.15537313, -0.26021054, + 0.12598538, -0.06024593, -0.11726313, -0.20979708, 0.20502815, + 0.22255987, 0.30414093, -0.11271469, 0.4672999 , 0.00509386, + -0.13042028, 0.22407965, -0.10005572, -0.14989175, -0.4365904 , + 0.2107629 , 0.18609835, -0.72566986, 0.18099388, -0.30312923, + 0.5998623 , -0.15350953, 0.36957732, -0.70576656, -0.68291605, + -0.420833 , -0.34906906, 0.06481054, 0.58362794, 0.31369215, + -0.82122093, -0.25893688, -0.24382648, 0.19694354, 0.16259049, + -0.11133604, -0.0934663 , -0.02774331, 0.21941422, -0.01251643, + 0.12317675, 0.6149345 , 0.14704953, 0.37044725, -0.28755513, + 0.1573624 , -0.44036826, 0.1894717 , 0.12188441, -0.49551368, + 0.14024805, -0.03632037, -0.11464448, -0.6433361 , 0.15439317, + -0.32633895, 0.00914915, -0.05114478, -0.1278209 , 0.33733654, + -0.07761955, -0.3304572 , -0.59981114, 0.14654088, 0.06847122, + -0.259691 , 0.25001425, 0.05990116, -0.645737 , 0.02534364, + -0.6641346 , -0.17594564, -0.01261867, 0.5100829 , 0.06452097, + -0.5573339 , 0.13375577, 0.25303102, 0.43233076, 0.15803403, + 0.254334 , 0.03359455, -0.05677273, -0.11994728, -0.2779966 , + -0.23615459, -0.04638248, -0.03050251, 0.3847538 , 0.0142238 , + -0.26559195, -0.03787645, 0.02767896, -0.08072381, -0.01213382, + 0.15982705, -0.07372479, 0.6117539 ], dtype=float32), + kernel: Array([[-0.21102534, 0.13572834, 0.08487151, ..., -0.37551054, + -0.08974288, -0.14673388], + [-0.42533755, 0.19477351, -0.05581117, ..., -0.19027376, + 0.11326433, 0.280313 ], + [ 0.20709127, -0.02206374, 0.00242551, ..., -0.21323436, + -0.1782914 , -0.49215955], + ..., + [-0.35879493, -0.13843112, 0.20357771, ..., -0.3228759 , + -0.48138356, -0.60722345], + [ 0.409117 , 0.09512895, -0.07367483, ..., -0.79642665, + -0.3589279 , 0.36061507], + [-0.1480074 , 0.24665533, -0.10206775, ..., 0.03880555, + -0.12229352, 0.42183685]], dtype=float32), + }, + value: { + bias: Array([-0.06382484, -0.10177656, -0.01783326, 0.07238453, 0.0269923 , + 0.03472366, 0.00803764, 0.09564783, -0.02505345, 0.16271482, + 0.01969427, -0.08568954, -0.02066216, -0.03939519, -0.09977195, + -0.01792192, 0.02864774, 0.03915608, 0.02681843, 0.07789411, + -0.12741898, 0.03238543, 0.01603028, -0.11089758, -0.08953978, + 0.11356429, -0.03502819, -0.0710902 , -0.07928857, -0.07034726, + -0.00949693, -0.11724807, 0.03667774, -0.0831453 , -0.08759701, + -0.11416655, 0.05479722, 0.02118164, 0.0808093 , -0.04038942, + -0.05000678, -0.06270798, -0.00748655, -0.0456011 , 0.09870113, + 0.03644619, 0.06150641, 0.15579787, 0.05361883, 0.07626446, + -0.02815951, 0.042865 , 0.06312688, 0.07170736, 0.185382 , + 0.06517194, 0.0584194 , -0.07202493, -0.09540485, -0.08599403, + -0.03540141, -0.10775873, 0.03609048, -0.06903887, -0.10958587, + -0.09116477, -0.10746475, -0.05029682, 0.07927579, 0.10405283, + 0.12020227, 0.03288461, -0.1196027 , 0.03575619, 0.0972458 , + 0.05244093, 0.08934712, 0.12486867, 0.08823898, 0.0983156 , + 0.10495333, -0.10362058, 0.09903006, 0.09088268, -0.12305731, + 0.10125972, -0.08239742, -0.00916524, -0.02156379, -0.09864786, + 0.10193784, -0.09551246, -0.05185878, -0.09967213, -0.10428255, + 0.10507633, -0.14955677, -0.12777895, -0.07334802, 0.16633663, + -0.04327126, 0.11767196, 0.08965011, 0.17071536, -0.0584787 , + -0.09368757, -0.16054668, -0.16359754, 0.0472418 , 0.1223535 , + -0.07023636, 0.24043013, 0.15963615, -0.01664593, 0.15456486, + 0.11985014, 0.25587445, 0.17214614, -0.15644138, 0.02051393, + -0.15587153, 0.03782795, 0.0117235 , 0.04015777, 0.17542166, + -0.1388207 , 0.09689201, -0.11741312], dtype=float32), + kernel: Array([[-0.06625638, -0.10676572, -0.11858384, ..., 0.12592293, + 0.19599795, -0.10273615], + [ 0.19458613, -0.08272579, -0.01620658, ..., 0.07205058, + 0.00837682, 0.03907323], + [ 0.17456186, -0.00289036, -0.24084125, ..., -0.03393131, + -0.21872245, 0.21209157], + ..., + [-0.10240208, -0.01977782, -0.14529926, ..., 0.48570195, + 0.29290506, -0.28927922], + [ 0.5828709 , 0.14644368, 0.23230499, ..., -0.51522744, + -0.00893875, 0.11504804], + [-0.18426676, -0.12852795, -0.02424787, ..., -0.17730054, + -0.39859936, 0.3791855 ]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-1.47393912e-01, 1.51925683e-01, -1.37644792e-02, -5.88501468e-02, + -2.74660420e-02, 2.99488921e-02, 6.36566132e-02, -3.54291759e-02, + -7.69458618e-03, 3.24805588e-01, -5.66644669e-02, 2.47724857e-02, + -1.12583041e-02, -9.56679955e-02, 1.36414379e-01, 3.80160026e-02, + 4.16892134e-02, -1.21877357e-01, -9.89358947e-02, 8.71492922e-02, + -5.38453227e-04, 6.59572333e-02, -4.53343103e-03, -4.99130785e-02, + 4.46245968e-02, 1.36791214e-01, -3.86863551e-03, 1.96640324e-02, + 6.48279041e-02, 6.96821213e-02, -2.48288596e-03, -8.72585997e-02, + 2.33484972e-02, -5.56293987e-02, 6.42355978e-02, -3.93086448e-02, + -5.08882850e-03, -8.52516666e-02, 1.37223052e-02, -1.85753062e-01, + 6.52511492e-02, -4.16962616e-03, -6.68387711e-01, -2.01383736e-02, + -1.15884997e-01, 1.01936243e-01, 4.85421829e-02, 2.81445622e-01, + 9.36387926e-02, -4.22131829e-02, -2.58780066e-02, -9.31250453e-02, + -8.50194786e-03, 1.86161455e-02, 1.08278200e-01, -1.33628231e-02, + -2.95319203e-02, -7.03121051e-02, 1.65721059e-01, 1.23157077e-01, + 1.26239538e-01, 8.69110152e-02, -6.91822842e-02, -6.16661599e-03, + -4.39871997e-02, 6.22018203e-02, -9.05643180e-02, 7.91532993e-02, + -6.81676865e-02, -5.36848418e-02, 2.38579530e-02, 2.05037221e-02, + 1.06868640e-01, -3.20243724e-02, 4.63792652e-01, 2.49073096e-02, + -1.17389895e-01, 3.09424363e-02, -7.55871981e-02, -5.58793545e-02, + -3.00298661e-01, 3.98461558e-02, -5.11545166e-02, -3.12989019e-02, + 6.69162199e-02, 7.48850405e-03, 8.32765177e-02, 1.84600934e-01, + 3.32334414e-02, 4.34078313e-02, 3.89995314e-02, 1.41597819e-02, + -9.58840549e-02, 1.53720230e-02, -5.27216345e-02, 6.26715794e-02, + -1.02486372e-01, -5.28356284e-02, -1.23486230e-02, 9.72601771e-02, + -4.58076522e-02, 1.89271364e-02, 6.57725483e-02, 7.59906247e-02, + -7.66134709e-02, -4.67620529e-02, 6.40749931e-02, 2.11690012e-02, + -8.14582258e-02, -7.82966167e-02, 3.28965578e-03, -1.84342973e-02, + 1.21189028e-01, -2.12921515e-01, -4.28700894e-02, -5.90704009e-02, + 1.14371911e-01, -1.21928066e-01, -3.76502164e-02, -3.38824801e-02, + -6.17372356e-02, -1.77653898e-02, 1.73333704e-01, 8.95729363e-02, + -6.20543491e-03, 2.95204688e-02, 2.87007019e-02, -4.14825082e-02], dtype=float32), + kernel: Array([[ 0.02049902, -0.1422533 , -0.4065625 , ..., 0.36780554, + 0.4333429 , 0.19315735], + [-0.03959692, 0.25549796, -0.06924722, ..., -0.55895936, + -0.27878842, 0.00849131], + [ 0.14230408, 0.490068 , -0.18883815, ..., -0.30614755, + -0.30285606, 0.07531235], + ..., + [ 0.0692085 , 0.07044961, 0.61252284, ..., -0.23293641, + -0.15725587, -0.13308391], + [ 0.19255584, -0.04722506, 0.23870187, ..., 0.26104337, + -0.1339734 , -0.28957063], + [-0.11599459, 0.07947025, -0.35309428, ..., -0.22296615, + 0.19523388, 0.4148677 ]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([-1.27752021e-01, -1.48293898e-01, -1.61956966e-01, 2.83768959e-03, + -1.43960312e-01, 4.62501980e-02, 1.55017406e-01, -1.05546750e-02, + 1.24534264e-01, 3.37903023e-01, -4.59039863e-03, -1.85061380e-01, + 1.73028678e-01, 2.01043487e-02, 2.81281001e-03, 9.39656422e-02, + -1.36595475e-03, -2.31539598e-03, -1.25028998e-01, 2.32474544e-04, + -2.05569100e-02, -1.70720860e-01, 9.77594405e-03, 1.18461205e-02, + 2.47885808e-01, -1.36246130e-01, -1.46172762e-01, -6.74006566e-02, + 1.35634257e-03, -6.88059255e-02, -1.38106868e-01, -3.34430300e-02, + 8.12808890e-03, 9.31725744e-03, 2.75329649e-02, -1.16437320e-02, + -1.76476855e-02, -1.26601487e-01, -3.02024838e-02, -5.67115694e-02, + 9.60008055e-02, 8.26168135e-02, -5.39703900e-03, 3.32480557e-02, + -1.42446086e-01, 2.15325981e-01, 3.59231502e-01, 6.70243725e-02, + 5.95594849e-03, 2.55793538e-02, 1.20438235e-02, -1.38192281e-01, + -9.54014882e-02, -6.51674941e-02, 1.10238351e-01, -9.56654698e-02, + 9.43305194e-02, -7.60812163e-02, -7.58638158e-02, 1.66560858e-01, + 1.86092202e-02, 1.56081989e-01, 5.11778006e-03, -4.85944413e-02, + 8.28391612e-02, 5.49042262e-02, -3.91565077e-02, -2.07471265e-03, + 3.06699127e-02, 3.01516578e-02, -3.39646004e-02, -2.84999385e-02, + 7.43485242e-02, 1.12809241e-02, -1.99580923e-01, -5.58756031e-02, + 2.26514554e-03, -5.92570044e-02, -5.33170663e-02, 4.89857085e-02, + 9.51478072e-03, -2.19953805e-02, -3.24405394e-02, -1.16561711e-01, + -7.09283957e-03, -4.68143495e-03, 5.30386604e-02, -1.35984283e-03, + 7.81681612e-02, -5.89879602e-03, -1.73093099e-03, -6.07467350e-03, + -3.00015986e-01, -3.05419937e-02, -1.54022276e-02, -6.91811442e-02, + -5.47736287e-02, 2.07600649e-02, 1.68445095e-01, -1.03444420e-02, + -1.04541652e-01, 5.39848953e-02, -2.24056616e-01, 1.65632993e-01, + -1.40168756e-01, -4.03944461e-04, -4.72264551e-02, 1.25392643e-03, + 6.72694370e-02, -2.77178794e-01, -4.07134034e-02, 8.86196047e-02, + 4.30512354e-02, -1.80928916e-01, -2.78730467e-02, 5.36768697e-02, + 2.03331217e-01, -2.78739282e-03, 1.01259798e-01, 6.90696239e-02, + -9.43698585e-02, 7.54763791e-03, 1.15448259e-01, 2.77357996e-01, + 1.10673435e-01, 1.47797644e-01, -9.77296499e-04, -1.15336455e-01], dtype=float32), + scale: Array([ 1.0143379 , 0.56084377, 0.61511797, 0.875825 , 0.7977206 , + 0.86066914, 0.4466555 , 0.5817147 , 0.8288987 , 1.2778934 , + 0.14663759, 0.92707175, 0.29109854, -0.5407728 , -0.00583301, + 0.80144596, 0.39370012, -0.01780795, 0.55510044, 0.8787668 , + 0.8118457 , 0.7212926 , 0.22898938, 0.00633279, 1.0344126 , + 0.4334817 , 1.2401807 , 0.45261577, 0.04603433, 0.5680823 , + 0.6459878 , 0.6229751 , 0.7747696 , 0.87726367, 0.4118688 , + 0.02244997, 0.7460163 , 0.53923315, 0.5421808 , 0.90724677, + 0.7915582 , 0.3890947 , 0.7099728 , 0.94361395, 0.44176197, + 0.708961 , 0.51787186, 1.1014241 , 0.3872233 , 0.6767877 , + 0.7677383 , 0.61901575, 0.73735505, 0.28867054, 0.81986 , + 0.68635523, 0.40406612, 0.7550665 , 0.3736026 , 0.42338955, + 0.51958174, 0.75038564, 0.10604075, 0.47208443, 0.53189516, + 0.43309155, 0.37580943, 0.48088777, 0.4441973 , 0.5321363 , + 0.5055118 , 0.98682684, 0.7596939 , 0.835741 , 0.10439195, + 0.57524043, 0.43323192, 0.71470016, 0.38151345, 0.8561705 , + 0.6086962 , 0.8476424 , 0.68529296, 0.81576556, 1.0155932 , + 0.2647536 , 0.66374 , 0.8123001 , 0.43478364, 0.7506543 , + 0.70008427, 0.01624053, 0.5848578 , 0.672363 , 0.5528674 , + 0.17722914, 0.7533496 , 0.50877273, 0.87119836, 0.50191104, + 0.70362985, 0.7272638 , 0.79176766, 1.0370609 , 0.6759593 , + -0.02001736, 0.8079056 , 0.76534927, 0.6989189 , 0.55686975, + 0.8860217 , 0.43745935, 0.6700682 , 0.90106094, 0.05214773, + 0.5131126 , 0.7131446 , 0.76999676, 0.6244328 , 0.87544256, + 0.54827476, 0.00163172, 0.72723454, 0.770018 , 0.62415475, + 0.49478632, 0.32370028, 0.52392954], dtype=float32), + }, + layernorm_before: { + bias: Array([ 5.46789803e-02, -1.76678717e-01, -3.22135426e-02, -1.92919653e-02, + -2.71175429e-03, -5.60209937e-02, 3.75961289e-02, 2.42887228e-03, + 2.20702612e-04, -1.47421882e-01, 1.25962566e-03, -9.49192513e-03, + 3.36418934e-02, 1.43754203e-02, 2.79071257e-02, 3.28558236e-02, + 1.25313839e-02, 5.57482056e-03, 2.42552720e-03, -4.17415686e-02, + -6.24468969e-03, -7.01998994e-02, 6.89587668e-02, -7.22786319e-03, + 1.22491186e-02, -1.73403710e-01, 4.65809219e-02, -2.40579285e-02, + 4.88756225e-04, 1.74035388e-03, 6.00601882e-02, -2.83498946e-03, + -4.87266015e-03, -6.68269256e-03, 2.96081766e-03, -4.06322069e-02, + -5.13824681e-03, -5.18148541e-02, -6.03306864e-04, 3.28535922e-02, + 5.79561479e-03, -5.60164126e-03, 2.76499480e-01, -5.52190188e-03, + 1.08755060e-01, 2.78750118e-02, -1.74138229e-03, -4.13134582e-02, + -1.12804314e-02, 1.25732794e-02, -2.00521722e-02, 3.98424082e-02, + 1.83249754e-03, -1.12617195e-01, 8.08609067e-04, 6.40430721e-03, + 4.37817909e-03, 2.31746268e-02, -2.95173544e-02, 5.21452923e-04, + 5.27834482e-02, 4.03346457e-02, -6.75705774e-03, -1.17142848e-03, + 2.27539521e-02, 8.85420591e-02, -6.18154998e-04, 2.34325184e-03, + 2.68001352e-02, 3.51031497e-03, -3.19339558e-02, 1.70021728e-02, + 5.05615212e-02, -3.64386290e-02, -5.41854799e-01, -3.14815044e-02, + -8.37918895e-04, -5.96139133e-02, -5.71792945e-03, 2.48509813e-02, + 5.07591176e-04, -9.80351865e-03, -4.27319445e-02, -4.39409167e-02, + -2.25920565e-02, 5.03404811e-03, 5.27535053e-03, -2.32479889e-02, + 8.92616529e-03, 2.46615335e-02, -1.04511445e-02, 4.72732633e-02, + 2.06987560e-03, -4.62427586e-02, 2.54220869e-02, -2.62153684e-03, + -1.93837471e-02, 3.46671464e-03, 9.97190922e-02, 5.05104400e-02, + -1.72052216e-02, 5.25105570e-04, -1.76775366e-01, 3.32803018e-02, + -2.44557224e-02, -1.54652279e-02, -1.59627572e-02, 3.77600156e-02, + -2.54226532e-02, -1.87684055e-02, -8.18012469e-03, 6.22503310e-02, + -9.38802143e-04, 4.75377999e-02, -9.10606887e-03, 2.95884646e-02, + 1.39154634e-02, -4.97860415e-03, 5.38583286e-03, 2.14410182e-02, + -7.47090429e-02, 2.04133317e-02, 1.69800371e-02, 5.70605770e-02, + -5.42162033e-03, -4.40403894e-02, 3.08437683e-02, -1.20276608e-01], dtype=float32), + scale: Array([ 3.09920963e-02, 1.43276036e-01, 1.49911180e-01, 2.15380549e-01, + 8.57652703e-05, 2.60982364e-01, 3.85197014e-01, -2.69828960e-02, + -3.57246934e-03, 1.77718669e-01, 1.78545862e-02, 2.83004701e-01, + 4.92948219e-02, 2.25244924e-01, 1.80071101e-01, 3.60851258e-01, + 2.13657439e-01, 6.50591101e-04, -1.27726868e-02, 8.39237347e-02, + 1.44251525e-01, 3.37201595e-01, 2.32133791e-01, 3.55153927e-03, + 1.30111605e-01, 3.19995940e-01, 1.58201158e-01, 1.93055466e-01, + 1.13361664e-02, 1.06982812e-02, 2.68632799e-01, 1.73866183e-01, + 2.49319479e-01, 2.33700909e-02, -3.93300736e-03, 1.54559046e-01, + 7.23777190e-02, 2.99003601e-01, 2.09477082e-01, 1.43370494e-01, + 2.27352500e-01, 2.16101669e-02, 1.23486273e-01, 3.68611127e-01, + 2.57967263e-01, 3.01627338e-01, 2.64851768e-02, 3.02257270e-01, + -8.79388452e-02, 1.15935393e-01, 3.20747286e-01, 3.84979308e-01, + 2.91317515e-02, 3.35988641e-01, 1.00166025e-03, -2.34178808e-02, + 1.65450841e-03, 9.69321802e-02, 1.97463453e-01, 2.78218836e-02, + 3.64540905e-01, 4.77716655e-01, -1.46322930e-02, 3.61578852e-01, + 2.67611772e-01, 1.56464770e-01, -1.12778105e-01, 3.55959013e-02, + 2.19024077e-01, 5.52458549e-03, 2.59881735e-01, 2.20885038e-01, + 2.75174618e-01, 2.36324608e-01, 1.31579578e-01, 2.25922078e-01, + 2.05405220e-01, 2.81007886e-01, -3.42874392e-03, 1.95928365e-01, + 1.01560242e-02, 1.41977131e-01, 2.52693713e-01, 2.46991277e-01, + 1.85481176e-01, -2.18331348e-02, -1.25332624e-02, 1.17323473e-01, + 6.74787387e-02, 1.07466944e-01, 1.33154258e-01, 1.68752119e-01, + -5.74604236e-03, 2.33412385e-01, 1.67403102e-01, -5.18426392e-03, + 1.57531217e-01, 1.38286993e-01, 3.68412942e-01, 2.31308982e-01, + 1.48976266e-01, 7.62487238e-04, 3.20091337e-01, 1.83057398e-01, + 2.31623769e-01, 2.69163251e-01, 1.83478698e-01, 3.07111442e-01, + 3.13205153e-01, 2.52505481e-01, 3.30679387e-01, 3.39435637e-01, + 9.45753418e-03, 4.02425140e-01, 6.86238613e-03, 2.50461936e-01, + 3.17101419e-01, 4.53172848e-02, 6.39474615e-02, 2.83503503e-01, + 1.75329089e-01, 1.88262612e-01, 1.82763785e-01, 3.37596089e-01, + 3.82400723e-03, 2.11784661e-01, 2.38452464e-01, 3.09545726e-01], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-3.9847344e-01, -3.6854276e-01, -4.5092517e-01, -4.3007284e-01, + -4.8035741e-01, 1.8460776e-03, -4.3567291e-01, -4.2967001e-01, + -3.3667952e-01, -4.1538629e-01, -3.6099383e-01, -4.3359479e-01, + -4.4380239e-01, -5.5524331e-01, -3.6870196e-01, -2.4764657e-01, + -4.7775912e-01, -4.1283187e-01, -3.4256911e-01, -2.0356853e-01, + -4.6999124e-01, -5.6414276e-01, -3.3079070e-01, -3.8981372e-01, + -4.5676839e-01, -8.1164964e-02, -2.3958890e-01, -4.1569802e-01, + -3.3400738e-01, -3.1094922e-02, -4.0657026e-01, -3.2647237e-01, + -5.5474406e-01, -5.1273751e-01, -4.6838763e-01, -4.5339519e-01, + -3.2953310e-01, -4.2982268e-01, -4.1318712e-01, -4.8618275e-01, + -3.3651689e-01, -3.8407576e-01, -3.7777162e-01, -4.4129571e-01, + -2.5866541e-01, -4.4932291e-01, -4.2343855e-01, -3.6280292e-01, + -3.8582063e-01, -3.7247223e-01, -7.4966386e-02, -3.5159770e-01, + -4.0223166e-01, -1.5560038e-01, -3.4897429e-01, -3.5570183e-01, + -2.3108204e-01, -3.0199406e-01, -5.0925249e-01, -4.2046854e-01, + -5.7922322e-01, -4.1019771e-01, -4.2622930e-01, -4.4688463e-01, + -4.6536162e-01, -3.6277786e-01, -2.4544826e-01, -3.0414537e-01, + -4.6335497e-01, -3.0622986e-01, -3.6244193e-01, -1.3396668e-01, + -4.9729797e-01, -4.0432706e-01, -1.7137310e-01, -1.5809147e-01, + -4.3105653e-01, -2.7127942e-01, -3.0011219e-01, -2.2747953e-01, + -3.7861067e-01, -2.8296992e-01, -4.9367398e-01, -1.4615911e-01, + -3.4952867e-01, -2.4669057e-01, -4.6957275e-01, -3.5128793e-01, + -3.9153334e-01, -4.2859459e-01, -4.0044680e-01, -4.3775627e-01, + -3.7979439e-01, -5.2443427e-01, -1.5687215e-01, -4.4120738e-01, + -4.7335675e-01, -3.9626312e-01, -3.6113608e-01, -8.2696810e-02, + -2.6075220e-01, -1.9168627e-01, -4.8281315e-01, -1.7972828e-01, + -5.2597886e-01, -2.7033496e-01, -5.4604769e-01, -3.4184673e-01, + -4.3216091e-01, -4.4862813e-01, -2.6511189e-01, -4.5229965e-01, + -3.8147661e-01, -3.2180479e-01, -3.1602433e-01, -4.3689847e-01, + -5.1146489e-01, -3.5368246e-01, -3.7338638e-01, -3.0413544e-01, + -2.4527813e-01, -4.7685048e-01, -2.5364643e-01, -3.0931181e-01, + -3.2056129e-01, -3.0224264e-01, -3.6615273e-01, -3.6344633e-01, + -4.7639742e-02, -4.2077422e-01, -4.8495576e-01, -3.2529086e-01, + -3.9314100e-01, -3.9539710e-01, -4.4750559e-01, -2.1847221e-01, + -3.4810659e-01, -3.5670453e-01, -4.8384792e-01, -4.6719027e-01, + -2.9682928e-01, -4.2099711e-01, -4.7452945e-01, -4.3598908e-01, + -4.5502350e-01, -4.7916898e-01, -5.1886410e-01, -2.9352316e-01, + -4.6964645e-01, -4.7157702e-01, -2.2555082e-01, -3.9752248e-01, + -7.7490084e-02, -1.6522288e-01, -4.9982330e-01, -4.2842680e-01, + -3.2544607e-01, -3.0850667e-01, -4.0218577e-01, -4.1208234e-01, + -4.0038344e-01, -3.7057555e-01, -2.2373217e-01, -3.4376973e-01, + -3.1466734e-01, -5.2123147e-01, -5.3340238e-01, -2.8239700e-01, + -3.3951777e-01, -4.2198521e-01, -4.4805121e-01, -2.5464025e-01, + -1.5028036e-01, -4.1943482e-01, -4.2202288e-01, -3.9689219e-01, + -3.7402752e-01, -1.6487882e-01, -3.0544364e-01, -4.0235436e-01, + -4.1456652e-01, -3.5549703e-01, -4.5845604e-01, -3.7517163e-01, + -4.1666454e-01, -2.5526059e-01, -4.9071726e-01, -3.1533489e-01, + -3.4268734e-01, -4.6997485e-01, -3.2526866e-01, -4.7221246e-01, + -3.2835072e-01, -5.4744834e-01, -4.1387323e-01, -2.6110759e-01, + -4.3684092e-01, -4.7436625e-01, -3.4099048e-01, -3.7735420e-01, + -4.5077652e-01, -5.1907909e-01, -3.7325633e-01, -2.2339475e-01, + -2.8594923e-01, -4.6022233e-01, -3.7742871e-01, -3.5255402e-01, + -1.6996916e-01, -3.3114195e-01, -2.1370201e-01, -4.9245229e-01, + -4.8635155e-01, -4.3522072e-01, -3.8971159e-01, -4.3678322e-01, + -4.6781665e-01, -2.9914606e-01, -4.2043447e-01, -2.9230228e-01, + -2.6290324e-01, -2.3281787e-01, -5.2061099e-01, -3.1569782e-01, + -5.0653201e-01, -4.7192627e-01, -3.8380232e-01, -4.1764271e-01, + -3.3932805e-01, -4.6573785e-01, -4.0508986e-01, -1.6529684e-01, + -3.9876980e-01, -4.5229152e-01, -2.8774762e-01, -3.5973960e-01, + -3.5486862e-01, -2.4070273e-01, -2.9540777e-01, -2.9138458e-01, + -3.0258924e-01, -3.0566293e-01, -3.6295635e-01, -2.6400667e-01, + -3.0755830e-01, -2.9665276e-01, -3.2570139e-01, -3.8138995e-01, + -2.5602007e-01, -3.7541685e-01, -4.3872544e-01, -4.1538125e-01, + -4.2229000e-01, -3.4041786e-01, -3.0963835e-01, -4.1050071e-01, + -5.6023502e-01, -5.6029701e-01, -3.5283986e-01, -2.6442611e-01, + -3.5243306e-01, -3.4913927e-01, -2.1650518e-01, -4.0154454e-01, + -3.7275749e-01, -4.7750309e-01, -3.6822703e-01, -4.5743442e-01, + -4.8160750e-01, -4.7312003e-01, -3.9876688e-01, -5.0516546e-01, + -5.4552037e-01, -3.4587166e-01, -1.6361688e-01, -3.0174541e-01, + -3.3792850e-01, -4.7918764e-01, -4.8411125e-01, -4.3018821e-01, + -4.6167120e-01, -4.1177577e-01, -4.5708910e-01, -4.5317286e-01, + -2.9856035e-01, -3.5552701e-01, -3.3068407e-01, -4.0677980e-01, + -4.2047119e-01, -2.2582465e-01, -5.4681146e-01, -3.3787546e-01, + -4.4846484e-01, -3.3120909e-01, -3.4642521e-01, -3.9533862e-01, + -3.4406123e-01, -4.5375788e-01, -5.0441438e-01, -3.3547118e-01, + -2.8920132e-01, -1.0944715e-01, -3.5370487e-01, -3.3076125e-01, + -2.8330564e-01, -3.6274710e-01, -2.7184573e-01, -3.5993487e-01, + -4.0513000e-01, -3.0320236e-01, -2.5303522e-01, -4.7531182e-01, + -4.6395004e-01, -4.2257047e-01, -4.8910224e-01, -3.6540136e-01, + 1.3740684e-01, -3.0950975e-01, -4.7773069e-01, -3.0934384e-01, + -4.1294369e-01, -4.5992157e-01, -4.8535463e-01, -4.4464013e-01, + -3.5180721e-01, -3.6469987e-01, -4.2818370e-01, -5.0174421e-01, + -3.8951018e-01, -3.7477994e-01, -3.4967715e-01, -4.3007034e-01, + -4.4603419e-01, -3.4457707e-01, -4.6689948e-01, -3.0864301e-01, + -4.9847844e-01, -3.7400711e-01, -4.7358903e-01, -3.5766226e-01, + -2.8434107e-01, -2.7686146e-01, -3.4817812e-01, -3.6055216e-01, + -4.3933406e-01, -5.2891338e-01, -2.7604723e-01, -5.1344776e-01, + -2.9326063e-01, -5.0666833e-01, -3.3879602e-01, -4.9309245e-01, + -3.0533355e-01, -5.3628761e-01, -4.3473902e-01, -4.1327205e-01, + -1.9567779e-01, -4.3977755e-01, -4.7691789e-01, -3.4455177e-01, + -4.1720390e-01, -3.4753558e-01, -4.0889883e-01, -4.2685688e-01, + -3.7366074e-01, -1.6097990e-01, -3.0323794e-01, -4.8915049e-01, + -3.5860199e-01, -3.2890522e-01, -3.6635777e-01, -3.0436921e-01, + -4.3726847e-01, -1.7787141e-01, -3.3810359e-01, -4.7570062e-01, + -5.1493889e-01, -4.9520645e-01, -4.8514831e-01, -4.0491089e-01, + -5.2531552e-01, -4.0086076e-01, -5.0189817e-01, -4.7050786e-01, + -2.1616998e-01, -3.5684329e-01, -4.2839581e-01, -3.8793033e-01, + -2.3242010e-01, -3.5336223e-01, -3.3516747e-01, -3.7301928e-01, + -4.4247723e-01, -4.8835176e-01, -2.0324853e-01, -3.9160007e-01, + -3.3982515e-01, -2.5195253e-01, -5.4986537e-01, -3.1654030e-01, + -4.6651766e-01, -4.7360784e-01, -2.8256953e-01, -3.3791304e-01, + -3.5749507e-01, -3.0531624e-01, -4.9014235e-01, -4.1924736e-01, + -2.0831045e-02, -3.5200366e-01, -2.2583617e-01, -4.5401311e-01, + -3.8651881e-01, -4.5588291e-01, -4.9305251e-01, -1.7810422e-03, + -2.6118660e-01, -3.1825066e-01, -2.4539576e-01, -2.5099796e-01, + -5.3763103e-01, -3.6526870e-02, -2.3981811e-01, -3.0151102e-01, + -4.7348249e-01, -2.2897437e-01, -4.1141531e-01, -4.8857701e-01, + -3.2772517e-01, -5.7252842e-01, -4.5862877e-01, -3.6546475e-01, + -3.6844012e-01, -2.0294371e-01, -4.8403230e-01, -4.8141652e-01, + -1.4463230e-03, -4.9992949e-01, -3.6349079e-01, -2.1532026e-01, + -4.1290241e-01, -4.8122567e-01, -2.8661153e-01, -4.4909769e-01, + -3.6796397e-01, -2.6207677e-01, -3.9517501e-01, -4.4621187e-01, + -3.6207435e-01, -4.8050162e-01, -3.7443826e-01, -2.3039897e-01, + -3.5041386e-01, -3.9968762e-01, -3.6701423e-01, -3.8556710e-01, + -3.4414786e-01, -4.8362854e-01, -3.8469809e-01, -4.8367801e-01, + -3.7582186e-01, -4.8235339e-01, -3.0144668e-01, -3.7805808e-01, + -3.9068311e-01, -4.6765688e-01, -4.5509905e-01, -4.1643125e-01, + -4.3739045e-01, -5.5417305e-01, -3.9493406e-01, -4.5592216e-01, + -3.2413527e-01, -4.4394735e-01, -2.5142363e-01, -2.9339787e-01, + -4.6713012e-01, -2.5261605e-01, -2.6347214e-01, -3.7425596e-01, + -2.1422139e-01, -4.4427580e-01, -3.5874227e-01, -4.1653034e-01, + -3.3887890e-01, -3.0552921e-01, -2.4492532e-01, -4.0205300e-01, + -3.2229105e-01, -4.8967350e-01, -2.1930854e-04, -3.7214684e-01, + -3.0033603e-01, -4.7503376e-01, -3.7032735e-01, -3.4375238e-01, + -3.8198787e-01, -3.3686790e-01, -5.5172253e-01, -3.0320832e-01, + -1.3836212e-01, -3.9621928e-01, -3.8446391e-01, -4.3740526e-01, + -3.8315192e-01, -4.0578344e-01, -3.2919899e-01, -4.7428545e-01, + -4.4342765e-01, -4.8948532e-01, -2.3972549e-01, -3.4704226e-01], dtype=float32), + kernel: Array([[ 0.10745826, 0.48113093, -0.04176594, ..., -0.36369812, + 0.15021859, -0.05807315], + [-0.58055264, -0.5796294 , -0.19744751, ..., -0.07468567, + -0.51942426, -0.03569252], + [ 0.01138834, 0.26800624, -0.01454838, ..., -0.8072465 , + 0.5447465 , 0.36949006], + ..., + [ 0.36528957, 0.27883154, -0.3109063 , ..., -0.5552432 , + 0.23742002, 0.05708292], + [-0.09702465, 0.6521357 , 0.16876654, ..., -0.252267 , + -0.09898252, -0.03783756], + [-0.05449405, -0.4971733 , -0.02605094, ..., 0.03078811, + 0.16231588, 0.4204838 ]], dtype=float32), + }, + output: { + bias: Array([-1.16826370e-01, 1.19446300e-01, -3.04981880e-03, -1.48974463e-01, + 5.59353642e-02, 7.24940300e-02, 3.09443232e-02, -3.97321768e-02, + -1.61944836e-01, -1.69336453e-01, -5.43983765e-02, 1.14368267e-01, + -1.82648357e-02, -5.84783889e-02, 2.02381879e-01, -1.69515777e-02, + 4.80070449e-02, -1.26781985e-01, -6.17002137e-02, 3.90345929e-03, + -1.18189584e-02, 1.06815845e-01, 2.18161605e-02, -5.02868928e-02, + -1.54914372e-02, 2.92950004e-01, 3.89679447e-02, 4.18758281e-02, + 6.82518259e-02, 8.17627087e-02, 2.73418054e-02, -7.99397752e-02, + 2.45839432e-02, 1.14206569e-02, 3.90483811e-02, -3.62712778e-02, + 2.89402343e-02, -7.93813691e-02, 1.43112317e-01, -1.52107775e-01, + 1.80989935e-03, -1.87356304e-02, -5.44708908e-01, -9.66665521e-03, + -5.63648678e-02, 1.56991724e-02, -2.02740412e-02, 5.39821014e-02, + 1.57445461e-01, -2.57236585e-02, -4.68497947e-02, -2.43687704e-02, + 9.51719135e-02, 7.29279891e-02, 1.93980023e-01, 1.36569049e-02, + 6.77157715e-02, -6.52143881e-02, 2.89444447e-01, 1.13720141e-01, + 1.87333062e-01, 1.36707956e-02, -7.90089220e-02, 1.85856335e-02, + -5.92513643e-02, 7.04728812e-02, 8.90799426e-03, 1.34663433e-01, + -6.09093085e-02, -2.86442898e-02, 3.39895934e-02, -6.57519400e-02, + -6.73585013e-02, -5.29800951e-02, 5.72383285e-01, 7.09643438e-02, + -1.11088164e-01, 3.76055352e-02, -4.64858226e-02, 1.60521493e-04, + -2.10771397e-01, -1.34013280e-01, -9.27246362e-02, -2.77775893e-04, + 1.43522322e-01, 4.07825224e-02, 5.58940060e-02, 1.22946128e-01, + -1.75536349e-01, 2.35335454e-01, 8.74865577e-02, 1.05813751e-02, + 1.50552183e-01, 9.60319340e-02, 6.48361864e-03, 6.70408085e-02, + -1.03584118e-01, -6.56064600e-02, -5.26508018e-02, 1.28201723e-01, + 3.76123935e-04, -2.51332112e-02, 8.39170814e-02, 3.41377631e-02, + -2.44789012e-02, -2.70240922e-02, 5.39324842e-02, -2.40320619e-02, + -9.49936584e-02, -1.08399363e-02, 8.49590078e-02, -4.56464440e-02, + -1.49793148e-01, 6.02403209e-02, -2.90313717e-02, -4.08411808e-02, + 2.29514278e-02, -1.53553560e-01, -5.11574633e-02, 3.26815434e-02, + 1.73677616e-02, -2.31411960e-02, 1.70279935e-01, -1.79555595e-01, + -7.74298981e-02, 9.48442426e-03, 4.52735312e-02, 9.07007884e-03], dtype=float32), + kernel: Array([[ 0.49678656, -0.04703587, 0.14592949, ..., 0.14652656, + 0.4102155 , 0.43552795], + [ 0.12093473, 0.28034988, 0.05580731, ..., -1.0471447 , + 0.00430787, -0.37116638], + [ 0.4650048 , 0.25468072, 0.04566461, ..., -0.30497208, + -0.02476864, 0.01837519], + ..., + [ 0.13562478, -0.04827686, 0.8059429 , ..., -0.31421518, + 0.24673045, -0.4773905 ], + [-0.73090583, -0.39090383, 0.1884505 , ..., 0.3404823 , + 0.5766392 , 0.37859675], + [-0.57174814, 0.14778472, -0.11844227, ..., -0.2646688 , + -0.18360172, -0.04887927]], dtype=float32), + }, + }, + }, + }, + }, + }, + layernorm: { + bias: Array([ 4.99482080e-02, -4.29396838e-01, 1.81105509e-02, 4.06332692e-05, + 1.38564035e-01, -2.67413799e-02, 4.30011488e-02, -3.71168442e-02, + -5.17193019e-01, -9.74547490e-03, -6.02933764e-02, 8.66996311e-03, + 3.20597798e-01, -1.68491941e-04, 1.20097917e-04, -8.34563673e-02, + 9.34075788e-02, -6.27656728e-02, -3.47385721e-05, -4.00960334e-02, + 1.16589768e-02, 2.07135342e-02, 5.25931716e-02, -5.63471504e-02, + 1.14345705e-04, -2.30510355e-04, 2.13737637e-02, -1.45103643e-03, + 1.10674717e-01, 3.27529460e-02, 1.30938485e-01, 2.17286926e-02, + -6.83023154e-06, -9.22439562e-04, -5.39473398e-03, -4.33185026e-02, + -9.40410607e-03, -1.07025966e-01, -1.36162853e-03, 1.07634941e-03, + 2.64154427e-04, -1.61221977e-02, 1.59604579e-01, 6.28793612e-04, + 2.51234043e-03, -1.61393564e-02, -3.47336605e-02, -2.62861013e-01, + -1.80254143e-03, 1.19618746e-02, -1.74341142e-01, -2.35029664e-02, + 1.11634769e-01, -7.43765607e-02, 2.52770900e-04, 5.80298016e-03, + 1.65572765e-05, -1.99696115e-05, 8.98374594e-07, 2.72183418e-02, + 5.92560282e-05, 2.25074757e-02, -7.86222667e-02, 3.12634259e-02, + -5.19049093e-02, 7.74146914e-02, -1.94988661e-02, 3.86598520e-04, + 3.62606696e-03, 1.13942362e-02, -5.74776307e-02, 1.42280420e-04, + -7.51663325e-03, -2.08205208e-01, -3.37208897e-01, 6.53023422e-02, + -9.40454900e-02, 2.59398657e-04, -1.21128708e-01, 6.60678595e-02, + 5.90708017e-01, -3.01386812e-04, -4.50036081e-04, -4.51468586e-05, + 1.67313181e-02, 1.51245928e-04, 1.40742473e-02, 5.68821088e-06, + 1.60795447e-04, 1.45528553e-04, 6.64302963e-04, 4.23284993e-02, + -2.40530949e-02, 2.40640657e-06, -1.34843877e-02, 4.14691959e-03, + -1.12099377e-02, -5.02210781e-02, 2.03482993e-02, 2.34702706e-01, + 7.58242968e-05, -5.44176670e-04, -2.74804551e-02, 2.38151588e-02, + -4.01919484e-02, -3.65443416e-02, -7.20889375e-06, -2.77027546e-04, + -8.81933421e-02, 1.69912830e-01, 2.06964463e-02, 1.79356318e-02, + -1.09324465e-04, 6.39357895e-04, 3.84408049e-02, 6.90272776e-04, + -3.20791632e-01, 6.30142495e-06, 4.83152457e-02, -2.25914060e-04, + -3.39954719e-02, -4.70114350e-02, 2.85215719e-05, 1.06935140e-05, + -1.23181306e-02, 1.01582222e-01, 3.10630426e-02, -4.47405800e-02], dtype=float32), + scale: Array([ 6.8868792e-01, 3.1249523e-01, 9.1828078e-01, -2.3013241e-04, + 3.7263566e-01, 2.3251632e-01, 6.5610683e-01, 1.1425257e-01, + 3.4526226e-01, 1.4145033e-02, 7.2997391e-01, 8.0623746e-02, + 2.9320392e-01, -4.1230087e-05, -6.6874339e-04, 5.7173258e-01, + 6.9446439e-01, 1.2802645e+00, -6.2002582e-05, 1.1703534e-01, + 4.1233909e-01, 8.1602544e-01, 5.0138348e-01, 5.7336593e-01, + 7.1710587e-05, 1.3667290e-05, 1.9022560e-01, 6.5288413e-01, + 4.8426083e-01, 4.8938426e-01, 6.1096561e-01, 4.0365598e-01, + 2.9291890e-05, 1.8095745e-03, 4.2584249e-01, 1.0999548e+00, + 2.1008846e-01, 3.8091195e-01, 4.8326021e-03, -1.2936177e-04, + -3.4556305e-04, 6.0092252e-01, 3.3731986e-02, 2.1269962e-03, + 7.6969045e-01, 3.1354058e-01, 5.6419522e-01, 3.5723236e-01, + 1.6803181e-01, 1.7196296e-01, 6.4958000e-01, 5.5659479e-01, + 2.6263377e-01, 2.4760671e-01, -5.8708456e-04, 3.1410387e-01, + -1.7643184e-04, 2.0361553e-04, -2.0005717e-05, 6.7649913e-01, + -1.2435010e-04, 3.2834163e-01, 6.4273733e-01, 4.8385906e-01, + 8.3392429e-01, 2.8222600e-01, 1.5083265e-01, -5.6320438e-03, + 3.3680016e-01, 2.1632074e-01, 4.0965131e-01, 4.5164741e-04, + 1.3090445e-01, 7.1241659e-01, 8.9404747e-02, 4.2915529e-01, + 6.3035327e-01, -3.9689415e-04, 6.1234730e-01, 2.9691941e-01, + 1.6158453e-01, 1.7273755e-04, 1.4721676e-04, -1.1007604e-04, + 1.3981619e-01, 1.9772084e-04, 4.9367583e-01, 2.4138608e-06, + -8.6720851e-05, 3.0593758e-03, 2.4073297e-01, 1.1165603e+00, + 1.5985191e-01, 3.0023075e-04, 2.5869986e-01, 9.9650544e-01, + 2.4187727e-01, 4.2451209e-01, 4.7070310e-01, 4.9373725e-01, + -3.1352194e-04, 1.4548376e-04, 1.2102582e+00, 2.8354323e-01, + 4.2492944e-01, 3.4629729e-01, 6.0579823e-06, -1.6931001e-03, + 7.3870343e-01, 6.7030573e-01, 1.5534639e-01, 5.8719957e-01, + -1.9493036e-04, 1.6623877e-03, 7.6086062e-01, 3.6656085e-01, + 3.9590457e-01, 7.1911036e-06, 7.4346775e-01, -2.0176137e-04, + 2.7498129e-01, 6.6088724e-01, -2.2207441e-06, -7.9591937e-06, + 2.0538671e-01, 1.2334855e+00, 2.2934711e-01, 1.6446222e-01], dtype=float32), + }, + }, +}) +FrozenDict({ + classifier: { + bias: Array([ 0.13959047, -0.15330061, 0.23963794, 0.26991433, 0.14193657, + -0.07303967, -0.15755221, -0.00429095, -0.1597117 , 0.01030699], dtype=float32), + kernel: Array([[-0.21733038, -0.04500001, 0.6502397 , ..., -0.16637883, + -0.19785312, -0.19238858], + [-0.05325412, 0.08365323, 0.04513923, ..., 0.04858669, + -0.2313862 , 0.15936573], + [ 0.15941077, -0.5403379 , 0.17727152, ..., 0.23605543, + -0.5477616 , 0.400377 ], + ..., + [ 0.14339526, -0.40216786, -0.11720804, ..., -0.4565188 , + 0.55370927, -0.3628201 ], + [ 0.1330832 , -0.05788376, 0.05994943, ..., -0.05280861, + -0.35830075, -0.07267395], + [ 0.07371017, 0.0797717 , 0.05526711, ..., -0.15986513, + -0.19333237, -0.11412273]], dtype=float32), + }, + vit: { + embeddings: { + cls_token: Array([[[ 4.21749912e-02, -7.63344541e-02, -9.38936546e-02, + 3.23799951e-03, 1.69068575e-01, 2.25809757e-02, + 3.44559725e-04, -3.61712719e-03, -8.03779662e-02, + 3.50343175e-02, -2.25252569e-01, 8.41705650e-02, + -1.40322876e+00, -6.44051209e-02, 3.56393009e-02, + 1.55636281e-01, 1.44712580e-02, -1.67391926e-01, + 1.13313071e-01, 2.03049585e-01, -1.21489443e-01, + 3.01832557e-01, -6.39758036e-02, 2.43316695e-01, + 8.19623619e-02, 3.64291817e-01, 2.20934585e-01, + 4.44887951e-02, -1.20811634e-01, 7.45030791e-02, + 2.88756996e-01, 2.45113447e-01, -1.02346905e-01, + -1.44495424e-02, -2.92423535e-02, 1.27275869e-01, + -4.68343794e-02, 6.18743040e-02, 1.62268832e-01, + 7.68043697e-02, 8.22397470e-02, 4.15349565e-02, + 4.29975651e-02, 3.34680080e-01, 1.23265073e-01, + -1.82213470e-01, -2.61384398e-01, -2.86995154e-02, + 4.78795953e-02, -7.12198839e-02, -9.28327814e-02, + -1.04730994e-01, -2.08872005e-01, 5.08915447e-02, + -2.92177856e-01, -4.63173315e-02, -1.39276087e-01, + -6.45741746e-02, 4.31298822e-01, -2.79950589e-01, + -3.42224717e-01, -2.57990032e-01, 7.88487419e-02, + 1.75719306e-01, -1.91833928e-01, -1.60736710e-01, + 2.24118829e-01, -1.42223332e-02, -1.56884298e-01, + -1.34610429e-01, -7.02320859e-02, -5.77486083e-02, + -2.67209888e-01, -1.95046023e-01, 1.72840059e-01, + 5.77186942e-02, 1.49908558e-01, 1.20706670e-01, + 2.60736823e-01, -4.85861525e-02, 1.50602728e-01, + -1.42758638e-01, 2.74297837e-02, 2.43728027e-01, + 2.83344448e-01, -3.49160563e-03, -1.21281445e-01, + -1.72878504e-01, -1.20296307e-01, -5.02357781e-02, + 8.33122134e-02, -1.96546650e+00, 1.50719002e-01, + 2.28078086e-02, -1.70492530e-01, 4.20110017e-01, + -9.01997983e-02, -3.95133868e-02, 1.77724794e-01, + 1.20884374e-01, -3.86571921e-02, -7.12739974e-02, + 1.50522381e-01, 1.34570688e-01, 3.98968756e-01, + -1.69036195e-01, 7.98634291e-02, -2.08218116e-02, + -1.60120368e-01, -2.48059444e-02, -1.43607566e-02, + 3.75894368e-01, 8.90320763e-02, -1.83609471e-01, + 3.58941674e-01, -1.86423570e-01, 2.57358737e-02, + 1.14504859e-01, -1.51107222e-01, -1.81861663e+00, + -1.02884192e-02, -3.45049873e-02, 7.34117031e-02, + 6.87620267e-02, 6.22403771e-02, 3.00719798e-01, + 6.51065707e-02, 3.67412001e-01]]], dtype=float32), + patch_embeddings: { + projection: { + bias: Array([-0.23706354, -1.6540319 , -0.56307954, 0.26430997, 0.45925793, + 0.10770348, -0.64220166, -0.256316 , -0.70986354, 0.00280434, + -0.8621466 , 0.33027068, -0.09428401, -0.18285578, 0.48031166, + 0.31876177, 0.1595892 , -0.7937292 , 0.41048172, 0.7553594 , + -0.11688971, 1.0029606 , -0.19344132, 0.9868722 , 0.33903363, + 1.4018133 , 0.6315837 , 0.15693058, -0.43901056, -0.03160467, + 1.4104275 , 0.8133011 , -0.8685446 , 0.09167413, -0.38324404, + 0.4127237 , -0.0992351 , -0.8858117 , 0.49157575, 1.1184502 , + 0.29291674, 0.09339268, 0.10110624, 1.3492149 , 0.23373768, + -0.62969214, -1.1953127 , -0.32950965, 0.16054444, -0.04147059, + -0.32538444, -0.44066548, -0.81780577, 0.3041742 , -0.88412094, + -0.2175792 , -0.7703008 , -0.29075882, 1.4584533 , -0.90483326, + -1.0301116 , -1.9739118 , 0.30468434, 0.47119713, -2.507297 , + -0.4836628 , 0.7882562 , 0.13242652, -0.64575255, -0.5046555 , + -0.2928836 , -0.19443981, -1.8653952 , -0.6447174 , 0.53096473, + 0.53206563, 0.53241354, 0.5417067 , 0.87263006, -0.22854666, + 0.57786745, -0.3563041 , -0.03160271, 0.5421835 , 1.0412699 , + -0.1696869 , -0.4523725 , -1.6923577 , -0.46177727, -0.20988323, + 0.2240901 , -0.9807764 , 0.61122894, -1.4953612 , -0.6922171 , + 1.6215447 , -0.32785478, -0.6734319 , 0.5917648 , 0.33969748, + -0.27886963, -0.50457156, 0.5443144 , 0.5749297 , 2.0102262 , + 0.15956564, 0.03995752, 0.18428732, 0.1132739 , 1.1407704 , + 0.10695937, 0.6835315 , 0.15768398, -0.5695771 , 0.45638716, + -0.6475178 , -0.24609064, 0.36820015, -0.7532235 , -0.60841304, + -0.13496049, -0.49064922, 0.31223625, 0.27192888, 0.27554023, + 1.1431948 , 0.30906013, 1.071581 ], dtype=float32), + kernel: Array([[[[ 0.1773554 , -0.0076261 , 0.39377195, ..., 0.09380915, + -0.01814741, 0.2283536 ], + [-0.2979896 , 0.02483123, 0.03855339, ..., 0.09292422, + 0.0840548 , 0.02128509], + [-0.11292487, -0.09941795, -0.21962844, ..., -0.00123785, + 0.14886236, 0.08236729]], + + [[ 0.32237047, -0.10110885, 0.5543121 , ..., 0.20536377, + -0.33312383, 0.189913 ], + [-0.23798265, -0.02816244, 0.15939426, ..., 0.1943017 , + -0.09701379, -0.03432931], + [ 0.10851432, -0.19066858, 0.04841118, ..., -0.01213947, + -0.04722768, 0.09524001]], + + [[ 0.2550473 , -0.04364953, -0.47635067, ..., 0.23234586, + -0.365624 , -0.01943434], + [-0.2977616 , 0.00419303, -0.8656642 , ..., 0.30416554, + -0.04954646, -0.1717201 ], + [ 0.11049867, -0.17219873, -0.6994334 , ..., 0.08241771, + 0.03823497, 0.0737647 ]], + + [[ 0.32358578, 0.17112552, 0.14320694, ..., -0.16366176, + -0.2593694 , 0.07770303], + [-0.1454706 , 0.12558867, -0.0796358 , ..., -0.1381298 , + -0.02053097, 0.09441353], + [ 0.05893492, -0.00259268, 0.02598703, ..., -0.20935844, + 0.1816749 , 0.30294153]]], + + + [[[ 0.14045075, -0.10724663, 0.1179781 , ..., 0.03702212, + 0.30327222, 0.0930538 ], + [-0.30190122, 0.11180247, -0.158114 , ..., 0.0602173 , + 0.28998795, -0.29604238], + [ 0.15810204, -0.02881167, -0.5261955 , ..., 0.16892877, + -0.02681437, -0.18120374]], + + [[ 0.1801053 , -0.17351657, 0.55432343, ..., -0.04310204, + 0.20605011, 0.15464272], + [-0.35629642, 0.1292508 , 0.4196015 , ..., -0.00466895, + 0.40400398, -0.22713828], + [ 0.3233468 , -0.09626149, 0.11579885, ..., 0.07394816, + -0.05250774, 0.03030762]], + + [[ 0.00362034, -0.27520794, -0.2702903 , ..., -0.0462058 , + 0.11103056, -0.17394325], + [-0.5339498 , 0.0243352 , -0.31235537, ..., 0.11845686, + 0.4861063 , -0.42657182], + [ 0.18858455, -0.18266024, -0.27887428, ..., 0.160679 , + 0.1009312 , -0.00895783]], + + [[ 0.23181222, -0.16900635, 0.17172067, ..., -0.2783022 , + 0.0531309 , -0.14735658], + [-0.17012398, 0.0160596 , 0.26340246, ..., -0.15997212, + 0.34474295, -0.18495837], + [ 0.25120145, -0.10030471, 0.4254995 , ..., -0.14183548, + 0.27963412, 0.15028709]]], + + + [[[ 0.04808923, -0.0793573 , -0.05729656, ..., -0.06481453, + -0.08433866, 0.25621825], + [-0.21680164, 0.12764114, -0.13568583, ..., -0.07495075, + -0.22111697, -0.15526403], + [ 0.3291076 , -0.05329537, -0.51923084, ..., 0.05456135, + -0.5494265 , -0.08629624]], + + [[ 0.00642845, -0.02851976, 0.48139414, ..., -0.3939782 , + -0.02460362, 0.4114907 ], + [-0.33011124, 0.21700603, 0.6363917 , ..., -0.34058657, + 0.07507597, -0.00235403], + [ 0.4467702 , -0.08816491, 0.338514 , ..., -0.20894945, + -0.44596782, 0.22598238]], + + [[-0.12848432, -0.17561927, -0.29103017, ..., -0.3294456 , + 0.0409772 , 0.05005151], + [-0.47656342, 0.09599458, 0.01909353, ..., -0.15514864, + 0.3270731 , -0.2028388 ], + [ 0.28001264, -0.17020059, 0.09817386, ..., -0.0626152 , + -0.07981837, 0.1589364 ]], + + [[ 0.18718387, -0.22170076, -0.12480265, ..., -0.2592998 , + -0.3301834 , 0.06356978], + [-0.05257114, -0.01859191, 0.32042652, ..., -0.23058891, + -0.06912984, 0.0211558 ], + [ 0.38780808, -0.15261608, 0.6172939 , ..., -0.2062019 , + -0.08611296, 0.31126183]]], + + + [[[-0.17952503, -0.05265503, -0.25038588, ..., 0.2504428 , + -0.02090847, 0.28380206], + [-0.2667295 , 0.07449401, -0.06040205, ..., 0.07079565, + -0.10550356, 0.09134477], + [ 0.09433349, -0.00351469, -0.22650217, ..., 0.02396235, + -0.11385609, -0.00377663]], + + [[-0.0575775 , 0.03152074, -0.16277562, ..., 0.3681371 , + 0.01227211, 0.08380491], + [-0.17844254, 0.10643482, 0.21974468, ..., 0.26805398, + 0.11603563, -0.11773686], + [ 0.38371345, -0.10176778, 0.20368902, ..., 0.1509468 , + -0.02080609, -0.06920536]], + + [[-0.18148401, -0.06984355, -0.85413545, ..., 0.47789526, + 0.04577804, -0.23575552], + [-0.3161512 , 0.02919929, -0.4008322 , ..., 0.44068208, + 0.26976907, -0.33418196], + [ 0.20420933, -0.1428987 , -0.11891793, ..., 0.30973166, + 0.22936389, -0.15018184]], + + [[-0.08356225, -0.13858244, -0.32757443, ..., 0.142595 , + -0.351965 , -0.0224065 ], + [-0.17430148, -0.03812229, 0.26883033, ..., -0.06516638, + -0.17817669, 0.04958347], + [ 0.09899019, -0.05975993, 0.6693031 , ..., -0.12235346, + 0.04438766, 0.1941795 ]]]], dtype=float32), + }, + }, + position_embeddings: Array([[[ 0.033918 , -0.0929377 , -0.10530281, ..., 0.28548247, + 0.08699364, 0.3533416 ], + [-1.1857578 , -0.05044684, 0.04907949, ..., 0.22999725, + 0.11194737, -0.5982713 ], + [-0.8083621 , -0.45285803, -0.35289305, ..., 0.1278877 , + -0.05470289, -0.3990702 ], + ..., + [ 0.01943148, -0.21790016, 0.20034477, ..., 0.13447031, + -0.12896743, 0.67419606], + [-0.5797325 , 0.09393977, -0.10601121, ..., 0.5656874 , + -0.4588539 , 0.3854702 ], + [-0.5565649 , 0.41388488, 0.21124075, ..., 0.8872757 , + -0.52099276, -0.30215293]]], dtype=float32), + }, + encoder: { + layer: { + 0: { + attention: { + attention: { + key: { + bias: Array([-0.03415372, 0.45620176, -0.6126655 , -0.51253253, 0.4352034 , + 0.649004 , 0.68406653, -0.46161497, -0.42964068, 0.8848562 , + 0.61135024, -0.67950135, -0.08493523, -0.29988992, 0.17360933, + 0.6337579 , 0.08232927, -0.59048873, -0.31515497, 0.09602205, + 0.14608672, -0.06844742, 0.5664715 , -0.6412251 , 0.44181392, + -0.40813568, 0.6399837 , 0.41530073, -0.66080755, -0.01420364, + -0.19360113, 0.14690916, -0.9573473 , -0.3561616 , -0.18794109, + 0.5325726 , -0.5846602 , -0.73184437, -0.35172847, -0.27912554, + 0.33543518, -0.21184462, 0.35979107, -0.12768705, 0.3412426 , + 0.35638675, 0.02133832, -0.08022975, -0.15245748, -0.20536591, + -0.11586879, -1.1425315 , -0.02779106, 0.32926312, -0.26967174, + 0.1869237 , -0.04547051, -1.0505885 , 0.9415589 , -0.5067646 , + -0.4434777 , 0.31517163, -0.48155847, 0.53517747, -0.40972438, + 0.16880448, 0.17367116, 0.29588723, 0.05554008, -0.48776305, + -0.31501368, -0.10502674, 0.13945673, 0.55460685, 0.51986885, + -0.39894032, -0.2494641 , 0.05012343, -0.04834368, 0.29485628, + 0.19806208, 0.3374114 , 0.11331975, 0.36502162, -0.01847783, + -0.39011198, 0.3790581 , 0.02779981, 0.651827 , 0.46813282, + 0.29877353, -0.43409863, 0.15295495, -0.35798907, -0.43806 , + -0.26218414, -0.45440933, 0.11940431, -0.4274145 , 0.18455578, + 0.00287167, 0.33210322, 0.32568368, -0.33072832, 0.08617178, + -0.24881381, -0.07380535, 0.23978767, -0.12747152, -0.7031561 , + 0.10753731, 0.27640837, 0.0658379 , -0.00729593, 0.15286936, + 0.14529614, 0.20602514, -0.18859431, 0.08264172, 0.09762905, + -0.1777643 , 0.21521233, 0.68952733, 0.35572144, -0.25402895, + -0.12086129, 0.33329782, -0.27421075], dtype=float32), + kernel: Array([[-0.13037314, -0.3519865 , -0.7311013 , ..., -0.14157563, + -0.098704 , 0.08183759], + [ 0.36583012, 0.46603498, 0.07223602, ..., 0.3872796 , + -0.01522054, 0.1815025 ], + [ 0.4379942 , -0.2856487 , -0.14937577, ..., 0.2555367 , + 0.143062 , -0.10283173], + ..., + [-0.08320192, -0.2315665 , 0.183584 , ..., -0.08580804, + -0.13162506, 0.07868645], + [ 0.42390886, -0.22044583, 0.4916073 , ..., 0.57441944, + -0.04296773, 0.02346938], + [ 0.09392057, -0.64890575, -0.21604928, ..., 0.09232247, + 0.13443615, 0.11167435]], dtype=float32), + }, + query: { + bias: Array([-0.06905509, -0.44198722, 0.30581513, 0.14314957, -0.15857255, + 0.15910107, -0.14257118, 0.11895534, -0.13231936, -0.28292307, + -0.1729494 , 0.1297792 , 0.29567292, 0.1645278 , -0.07054784, + -0.1081412 , 0.2331662 , -0.02096415, 0.00768806, -0.13316594, + -0.25273684, -0.09793813, 0.04979343, 0.23559354, -0.19012289, + 0.08803113, -0.03696692, -0.29770622, 0.39657098, 0.1723672 , + 0.12946679, 0.12689376, 0.3535053 , -0.02396139, 0.01831555, + -0.45482886, 0.0035749 , 0.0489781 , -0.19142908, 0.04616706, + -0.10055473, 0.00193944, -0.0123636 , -0.1297169 , 0.06566181, + -0.12390063, 0.01915232, -0.10427503, -0.26933163, 0.01553594, + -0.02893843, 0.49531803, 0.30729276, 0.01837071, 0.172419 , + -0.05943993, -0.29665995, 0.46334326, -0.13110274, 0.49487928, + 0.21740474, -0.25742924, 0.03184005, -0.16928338, 0.04680282, + -0.10523521, 0.07657482, 0.147745 , -0.09269993, 0.01065696, + 0.05368982, -0.18941729, 0.10891443, 0.20821781, 0.15085766, + 0.51667964, -0.15096618, 0.34021842, -0.03054576, 0.26622248, + -0.36788234, 0.29284194, -0.18616207, 0.09047277, 0.3589835 , + -0.09647398, -0.10208427, 0.09295312, 0.03126433, -0.27378327, + 0.08303533, -0.17912042, -0.06608632, -0.05499709, -0.09980894, + 0.15720649, 0.37134722, 0.39746282, 0.172588 , 0.04140972, + 0.58430547, -0.250469 , -0.3090034 , 0.35163316, -0.09326122, + 0.18299101, -0.02789885, -0.16457352, 0.21470276, 0.20549582, + -0.27671963, -0.4622744 , 0.19656613, 0.31474093, -0.0242331 , + -0.01637263, 0.08560522, 0.40495184, -0.14052592, -0.24214344, + -0.27252468, 0.15204628, -0.19363014, -0.13963734, -0.02372985, + -0.13916686, 0.21864863, 0.13198575], dtype=float32), + kernel: Array([[-0.02407378, -0.08052919, -0.78489935, ..., 0.0214601 , + -0.21603775, 0.07814599], + [ 0.1441562 , 0.39637634, -0.01425371, ..., 0.44336745, + 0.02438617, 0.04486706], + [ 0.54582417, -0.38129556, -0.09409017, ..., 0.07300002, + 0.08039478, -0.01730129], + ..., + [-0.40244198, 0.23577765, -0.09792052, ..., 0.06759664, + 0.0144011 , -0.21112448], + [ 0.52981406, 0.07420167, 0.320484 , ..., 0.3295426 , + -0.08885895, -0.11461516], + [ 0.2322218 , -0.26398233, -0.4191451 , ..., 0.24738975, + -0.05001807, 0.01805569]], dtype=float32), + }, + value: { + bias: Array([-2.94787940e-02, -2.72796545e-02, 6.14140369e-02, -2.86302455e-02, + 7.17062429e-02, -3.11819911e-02, -4.36154753e-02, 3.26400585e-02, + -2.82624215e-02, 1.26899080e-02, -7.74416178e-02, 4.88068024e-03, + 2.13063625e-03, -7.00014010e-02, 6.05225153e-02, -4.30300795e-02, + 1.79692768e-02, -2.71751042e-02, 3.53089422e-02, -1.95429940e-02, + -1.21456282e-02, -2.05201493e-03, -3.62878200e-05, -1.08139282e-02, + 7.40849450e-02, 1.50353080e-02, -7.06956685e-02, 8.49624351e-02, + -7.26162344e-02, -5.55856004e-02, -7.41069913e-02, 7.11530168e-03, + -2.33412739e-02, 6.24122694e-02, 1.00617625e-01, 4.78886962e-02, + 6.93868250e-02, 8.44263211e-02, 6.89995065e-02, -1.49511993e-02, + 6.76077977e-02, 2.99593173e-02, -6.50391206e-02, -3.10420450e-02, + -4.87265997e-02, -7.07542598e-02, 3.83699648e-02, -1.07949436e-01, + 3.88861708e-02, 1.07101694e-01, -7.05710873e-02, 6.53532194e-03, + 5.76992594e-02, 8.30659941e-02, -1.04475126e-01, 8.40004310e-02, + -9.89107862e-02, 8.94425064e-02, -3.12575921e-02, -5.92782833e-02, + -8.53339434e-02, -5.87267503e-02, 3.41176279e-02, -6.32282570e-02, + -4.08555977e-02, -7.24948272e-02, 4.16365713e-02, -2.20883206e-01, + -4.90566827e-02, -1.89475138e-02, -1.50290072e-01, 2.09986553e-01, + 3.20502929e-02, 8.28777160e-03, 6.13699071e-02, -2.50448077e-03, + -2.15128988e-01, -3.07804886e-02, -1.82573318e-01, 9.74187329e-02, + 8.86107534e-02, -6.78683072e-02, -6.89225420e-02, -3.00914403e-02, + -1.05735838e-01, -2.17407808e-01, 1.63544074e-01, 1.00356489e-01, + -6.03550719e-03, 3.53857167e-02, -7.28843063e-02, 1.15036063e-01, + -1.29502550e-01, 5.06461645e-03, 1.20808132e-01, 1.27421441e-02, + -3.18243355e-02, -5.57740144e-02, -6.90745190e-02, -1.28188869e-02, + 1.20238643e-02, -4.42795120e-02, -2.79449914e-02, -3.33387293e-02, + 7.23095536e-02, -2.58487016e-02, -1.62459761e-02, 1.07147522e-01, + -6.11629374e-02, 7.07705468e-02, -5.56649491e-02, 3.85964476e-02, + 9.34836396e-04, -1.59531720e-02, -7.20297545e-02, 8.05353299e-02, + 1.81769580e-02, -3.95291820e-02, -7.38437474e-02, -9.54268277e-02, + 1.34773731e-01, 9.29158255e-02, -5.70693836e-02, 2.48423647e-02, + -8.44348520e-02, -7.13706389e-02, -7.08770528e-02, 7.93647394e-02], dtype=float32), + kernel: Array([[-0.19559653, 0.10921917, -0.24855891, ..., -0.3095954 , + -0.11501747, 0.3495971 ], + [ 0.2759333 , -0.0968355 , 0.25864232, ..., -0.07446828, + -0.01653024, 0.04277551], + [-0.53325766, -0.17772451, -0.19101648, ..., -0.09837775, + 0.07115 , 0.39424804], + ..., + [ 0.36891153, -0.11530672, 0.24386573, ..., 0.05493925, + -0.2263004 , -0.10157143], + [ 0.06181217, 0.21771231, -0.52629954, ..., -0.0283559 , + -0.11670491, -0.19776788], + [-0.20297265, -0.24884045, 0.21022551, ..., 0.13422379, + 0.22147001, 0.18063086]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-0.08976492, 0.35623267, -0.28853896, 0.1584903 , -0.04931528, + 0.41930285, -0.12198956, -0.1676379 , -0.07629202, 0.44505563, + -0.13279814, -0.2606873 , -0.38957864, 0.2600788 , -0.15301487, + -0.7227959 , 0.06558075, -0.02327237, -0.01965788, 0.3900318 , + 0.14463983, 0.33909434, -0.02560525, 0.25102413, 0.26336563, + 0.07151037, -0.01791747, -0.32692116, -0.29931375, 0.26871207, + -0.45228213, -0.07384168, -0.2372598 , 0.16236512, 0.20462725, + -0.04373658, 0.18107706, 0.5327047 , 0.11542407, -0.20473395, + -0.29505363, -0.1296326 , -0.50376815, 0.23728877, 0.16542588, + 0.20167176, 0.05926127, 0.1154462 , -0.14579555, -0.08487356, + 0.07978057, -0.27115968, -0.5287936 , 0.07883674, -0.07367402, + -0.04930678, -0.5660592 , 0.14252122, -0.05342535, -0.03159878, + 0.18108739, -0.06234799, 0.13124788, 0.00128543, -0.24720064, + 0.13804607, -0.5026876 , 0.53043497, 0.07013435, 0.22524473, + 0.17431633, 0.22486666, 0.01429994, -0.02526889, 0.65786517, + -0.12362218, -0.124548 , 0.02234903, -0.16374543, -0.03648356, + -0.79906535, -0.21115243, -0.20314318, 0.24117418, -0.42399853, + 0.37275848, 0.11345251, -0.07611627, 0.09416302, -0.16610427, + 0.22861136, -0.22463125, 0.21705595, 0.09253941, -0.2043051 , + -0.06828043, -0.27955458, -0.00722787, -0.13780643, -0.44847608, + 0.18490234, 0.09589677, 0.2274974 , -0.04555204, -0.13882257, + -0.17222178, -0.24243575, 0.09366386, 0.16359329, 0.15961671, + 0.08576287, 0.09707799, 0.31846833, -0.27176753, 0.12482687, + -0.35437414, -0.04145558, -0.4648349 , 0.1269856 , -0.09939937, + 0.00209725, 0.00080232, 0.20968738, -0.18917999, -0.12978964, + 0.13607804, 0.00752292, -0.00256601], dtype=float32), + kernel: Array([[ 0.01726927, 0.12866902, -0.1439863 , ..., -0.26570368, + -0.0510007 , 0.47683123], + [ 0.15968257, -0.1014777 , -0.26105717, ..., 0.30918548, + 0.48335168, 0.01941068], + [-0.05697384, 0.05050297, -0.3797851 , ..., -0.70562077, + 0.11496508, 0.02426186], + ..., + [-0.37025422, -0.41289684, -0.09372032, ..., 0.3905355 , + -0.44608384, -0.10514601], + [ 0.03965069, -0.21772908, -0.06686966, ..., 0.5678079 , + -0.5982024 , 0.03439448], + [ 0.31467485, 0.03231593, 0.6761289 , ..., 0.3369738 , + -0.16577064, 0.09921929]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([ 7.65252188e-02, 6.29107952e-01, 8.02198052e-02, 8.70596059e-03, + 4.38350849e-02, -1.07495952e-02, -5.31466417e-02, -1.02623366e-01, + 2.12361347e-02, 6.80099875e-02, -2.28481810e-03, 8.53950828e-02, + -5.73644787e-02, 2.06265952e-02, -7.53715783e-02, -1.43298656e-01, + -2.74074376e-02, -1.09807447e-01, -6.17579147e-02, -7.90846869e-02, + -9.07771438e-02, 8.38840082e-02, -1.48849720e-02, -2.24489588e-02, + -9.19756591e-02, -1.33614585e-01, 8.40165392e-02, -2.04390988e-01, + 7.47799221e-03, 1.38857782e-01, -5.69530539e-02, 9.17964336e-03, + -1.55208111e-01, 2.73925383e-02, -1.95245608e-03, 5.41600175e-02, + -2.93773953e-02, 3.53689879e-01, -1.23494409e-01, -1.56066343e-01, + 8.71012285e-02, 2.55195685e-02, -3.10201854e-01, -8.39867163e-03, + -3.08605190e-02, -5.28353490e-02, -1.40786827e-01, -5.47004975e-02, + -1.13185709e-02, 1.86418276e-02, -7.84755684e-03, -4.15783050e-03, + -7.23032653e-02, -1.61183357e-01, 1.90100819e-02, -4.55883006e-03, + -2.70947274e-02, 1.69784650e-02, -2.28463719e-03, 4.21662033e-02, + -9.98298172e-03, 7.62146562e-02, -1.09199071e-02, 7.20275193e-02, + 1.14038222e-01, 2.85849470e-04, -2.99424119e-02, 1.37919337e-01, + -1.45379841e-01, 6.62748562e-03, 2.23827194e-02, -5.66348387e-03, + 5.01772054e-02, 1.41351158e-03, 5.96981421e-02, -2.08293200e-01, + -5.95040731e-02, 5.53276669e-03, -8.04201216e-02, -1.67969614e-01, + 3.61396512e-03, 5.54506434e-03, -2.92952433e-02, 4.65981402e-02, + 1.59924058e-03, -2.43443865e-02, -1.17824182e-01, -3.83222140e-02, + -3.87748890e-02, -6.59047961e-02, -8.08075536e-04, -3.35978210e-01, + -4.24195863e-02, 3.55881304e-01, 3.02827614e-03, -5.55533350e-01, + 1.54737878e-04, -1.47938222e-01, -2.13265121e-02, -1.14879206e-01, + -3.36717785e-04, 5.57862930e-02, 1.95077602e-02, -8.35356787e-02, + 2.89721284e-02, -1.48645073e-01, -2.94460828e-04, 8.76041055e-02, + -1.90158889e-01, 1.32130921e-01, 1.11547068e-01, -8.59021768e-02, + -4.04770300e-03, -2.04068184e-01, 2.25875482e-01, -7.98009858e-02, + 1.03451028e-01, -6.95376145e-03, -8.43686834e-02, -1.78507604e-02, + -9.83250737e-02, -7.10767433e-02, -6.19554799e-03, -5.51361814e-02, + -8.06146394e-03, 4.24989313e-01, -2.02743262e-02, -7.10928217e-02], dtype=float32), + scale: Array([ 4.92835701e-01, 7.25533962e-01, 1.97780564e-01, 7.19367981e-01, + 3.85366410e-01, 1.70830980e-01, 4.03196029e-02, 3.59184593e-01, + 1.64954796e-01, 5.99162340e-01, 1.94611892e-01, 4.47059900e-01, + 1.33249581e-01, 1.17341705e-01, 2.57417858e-01, 2.13419840e-01, + 1.17401130e-01, 1.44011736e-01, 4.55426246e-01, 3.37985694e-01, + 3.68331194e-01, 1.58645034e-01, 4.11303312e-01, 2.80974478e-01, + 2.08782077e-01, 3.16983044e-01, 2.54319280e-01, 5.07932603e-01, + 6.26261353e-01, 3.31194192e-01, 5.77102304e-01, 5.91307692e-03, + 2.17876673e-01, 1.83340341e-01, 7.28062959e-03, 1.43880785e-01, + 3.79934996e-01, 8.64668906e-01, 5.03845096e-01, 3.30553919e-01, + 3.66227090e-01, 1.23913087e-01, 4.04930204e-01, -5.25120134e-03, + 2.16952696e-01, 2.82419056e-01, 2.41874039e-01, 2.18594924e-01, + 2.72880584e-01, 1.63890943e-01, 5.69228344e-02, 1.57258492e-02, + 3.11221033e-01, 3.75806093e-01, 3.45861703e-01, 3.33030730e-01, + 1.86815023e-01, 2.87897736e-01, 1.79882497e-02, 3.26225400e-01, + 3.56188752e-02, 3.17295760e-01, 5.93220115e-01, 3.10718954e-01, + 1.92407712e-01, 9.49200336e-03, 2.15880141e-01, 4.15735185e-01, + 3.10675055e-01, 5.28817117e-01, 1.04173638e-01, 1.57576382e-01, + 3.02691996e-01, 3.79474498e-02, 7.07599940e-03, 5.08240998e-01, + 3.50080490e-01, 3.61189306e-01, 1.87464565e-01, 4.86818105e-01, + 3.34662758e-03, -1.36662647e-03, 1.20725386e-01, 5.04414022e-01, + 3.98086966e-04, 9.34311152e-02, 1.42500699e-01, 3.84391755e-01, + 1.07679166e-01, 3.01436037e-01, -1.05553423e-03, 4.41559106e-01, + 1.26928464e-01, 4.32251424e-01, -6.94631459e-03, 5.02396822e-01, + -3.20078339e-04, 1.94142312e-01, 8.77228081e-02, 1.18769050e-01, + 4.47774172e-01, 3.87618184e-01, 4.12855119e-01, 3.15822661e-01, + 4.97932345e-01, 6.33932471e-01, 6.15001982e-03, 2.09667221e-01, + 2.26229250e-01, 1.05634868e+00, 5.23372293e-01, 4.29101616e-01, + 2.27578685e-01, 6.58651769e-01, 2.97994912e-01, 1.07822992e-01, + 4.83628422e-01, 4.72095191e-01, 3.23185503e-01, 1.00680351e-01, + 4.04555202e-01, 1.64712802e-01, 6.41667992e-02, 1.81472510e-01, + 5.06515145e-01, 4.70199496e-01, 2.70896405e-01, 4.36434001e-01], dtype=float32), + }, + layernorm_before: { + bias: Array([-0.05890432, 0.07051883, 0.09445145, -0.07912993, -0.14302658, + -0.01429262, 0.03426607, 0.00171352, 0.11276561, -0.05662137, + 0.13313785, -0.04794319, 0.00357568, 0.01197018, -0.07405796, + -0.06490814, -0.0245945 , 0.14022568, -0.05081064, -0.12518148, + 0.00789185, -0.24005446, 0.02322139, -0.17626956, -0.03338344, + -0.33740795, -0.13774505, -0.07997898, 0.04349966, -0.01520896, + -0.39623648, -0.150468 , 0.11885046, -0.01181321, 0.00808169, + -0.20219389, 0.0043777 , -0.08942726, -0.16421181, -0.18974544, + -0.07742172, -0.01755824, -0.02627599, -0.4378331 , -0.19054742, + 0.2720697 , 0.25448796, 0.01773962, -0.01873611, 0.0169231 , + 0.03406947, 0.0349858 , 0.09072868, -0.05627641, 0.14739008, + 0.01965971, 0.09791738, 0.03601414, -0.28183195, 0.15014632, + 0.22542529, 0.3072184 , -0.09403451, -0.08351854, 0.25167593, + 0.04495538, -0.11564947, -0.00397131, 0.10662062, 0.05204451, + 0.01583207, 0.01414629, 0.23943776, 0.06893958, -0.08743864, + -0.06693614, -0.07541402, -0.19011301, -0.17251883, 0.01803088, + -0.09193094, 0.04894969, -0.00502513, -0.09108073, -0.1768744 , + -0.01100582, 0.10178719, 0.22854435, 0.04437195, 0.00859537, + -0.03013818, 0.26522186, -0.0835686 , 0.0909194 , 0.08486014, + -0.39586094, 0.03491055, 0.02060359, -0.14571564, -0.18893251, + 0.04928007, 0.05034845, -0.15555385, -0.15786135, -0.56710386, + 0.01834173, -0.04088113, -0.0075858 , 0.04828192, -0.1228609 , + -0.00563954, -0.3955512 , -0.06544469, 0.09136705, -0.10234332, + 0.09786685, 0.00981475, -0.05797027, 0.08497368, 0.18863663, + -0.00070434, 0.03551389, -0.04786862, -0.03968854, -0.02606042, + -0.2953443 , -0.03344315, -0.2232286 ], dtype=float32), + scale: Array([0.45335642, 0.28725365, 0.25888288, 0.51014036, 0.25393826, + 0.14649442, 0.45231742, 0.07753127, 0.3569898 , 0.67538893, + 0.2085993 , 0.17448373, 0.01005436, 0.07782084, 0.26946452, + 0.14676209, 0.22900493, 0.29659525, 0.13212137, 0.1925922 , + 0.04533805, 0.26069018, 0.15647988, 0.20724292, 0.08556735, + 0.296176 , 0.20237958, 0.6578139 , 0.11472861, 0.3509395 , + 0.38030067, 0.19699657, 0.30957213, 0.1079342 , 0.12253352, + 0.54774934, 0.06868377, 0.49437425, 0.34611228, 0.33752495, + 0.24485146, 0.11372902, 0.13047746, 0.40855378, 0.49101192, + 0.5738504 , 0.32902393, 0.2079778 , 0.09577054, 0.11233962, + 0.17488489, 0.10529607, 0.13811632, 0.19155163, 0.16904642, + 0.10060844, 0.2304541 , 0.20629783, 0.21513213, 0.1885996 , + 0.2502167 , 0.3126915 , 0.3413607 , 0.1449984 , 0.41228867, + 0.09932974, 0.14989796, 0.13950281, 0.24986736, 0.15920106, + 0.06112578, 0.10076807, 0.25724828, 0.11337321, 0.1927693 , + 0.24177718, 0.14134662, 0.4839469 , 0.20330684, 0.1124863 , + 0.17910473, 0.1306797 , 0.09654108, 0.13909315, 0.19018278, + 0.12305983, 0.3541507 , 0.31824607, 0.13942654, 0.08698463, + 0.08841392, 0.17479059, 0.15011418, 0.45322958, 0.17411204, + 0.29569057, 0.10512718, 0.3171436 , 0.21283199, 0.51158214, + 0.42643854, 0.18322334, 0.3302157 , 0.37901062, 0.44037428, + 0.50349325, 0.1318383 , 0.13000993, 0.14093749, 0.5986485 , + 0.1306848 , 0.36673996, 0.2426187 , 0.19750752, 0.14503735, + 0.22176953, 0.1925638 , 0.15475766, 0.17507745, 0.22627683, + 0.10930963, 0.37122646, 0.12475897, 0.12880948, 0.11971601, + 0.2903781 , 0.1446715 , 0.20750394], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-0.4722571 , -0.43928552, -0.5569621 , -0.3837599 , -0.47573015, + -0.4316669 , -0.762032 , -0.5649565 , -0.31526604, -0.42409667, + -0.35039356, -0.5466656 , -0.08653939, -0.50116634, -0.65188706, + -0.47724712, -0.29790252, 0.32602024, -0.04504937, -0.40849677, + -0.11292349, -0.25168103, -0.45442316, -0.34808317, -0.38236308, + -0.3788954 , -0.19183853, -0.30315536, -0.16824003, -0.4485127 , + -0.30100617, -0.39639288, -0.16540907, -0.5450078 , -0.4645357 , + -0.36592278, -0.5585429 , -0.34257036, -0.19063933, -0.44469953, + -0.42737913, -0.36188826, -0.2907807 , -0.02146477, -0.23154494, + -0.20253362, -0.69757545, -0.39887667, -0.44025323, -0.39642838, + -0.387052 , -0.24472845, -0.09177764, 0.05788963, -0.04718702, + -0.21287683, -0.29563203, -0.3253591 , -0.43122077, -0.1849469 , + -0.5141688 , -0.29721844, -0.47572163, -0.3446094 , -0.26447332, + -0.47591433, -0.19857107, -0.6629566 , -0.36350504, -0.60418856, + -0.4550368 , -0.40432703, -0.35896167, -0.33915696, -0.41339654, + -0.23302272, -0.05349919, -0.47071078, -0.32239798, -0.3542714 , + -0.25527537, -0.3346329 , -0.49991933, -0.22825265, -0.45629248, + 0.01691399, -0.34854415, -0.44188562, -0.03925477, -0.3642511 , + -0.40468022, -0.33773935, -0.49415597, -0.5592706 , -0.2977098 , + -0.32126102, -0.16516496, -0.40793797, -0.8022191 , -0.44280693, + -0.1008477 , -0.22637095, -0.16032524, -0.45757398, -0.4884214 , + -0.48628846, -0.4652781 , -0.51417524, -0.25536695, -0.33155087, + -0.5587205 , -0.37767413, -0.30986476, -0.47982594, -0.51307243, + -0.38967907, -0.24242653, -0.25831246, -0.1458375 , -0.58483136, + -0.3381596 , 0.27927315, -0.7453308 , -0.4533877 , -0.15844217, + -0.45824784, -0.468726 , -0.4619875 , -0.4371935 , 0.00825728, + -0.31616744, -0.4558102 , 0.0993996 , -0.08013049, -0.17536967, + -0.52361816, -0.44187388, -0.25537458, -0.19660369, -0.30101493, + -0.09137805, -0.41055298, -0.46738333, -0.57409215, -0.34877914, + -0.52286893, -0.55365384, -0.33294562, -0.47346646, -0.3105509 , + -0.44117743, -0.3502154 , -0.39879954, -0.3967916 , -0.47747064, + -0.24431925, -0.39359137, -0.22136432, -0.35870054, -0.29925522, + -0.34952432, -0.49339256, -0.42442355, -0.383239 , -0.3228079 , + -0.59250957, -0.3766234 , -0.48536432, -0.49891275, -0.10365842, + -0.21082799, -0.5129332 , -0.39634183, -0.4714169 , -0.3876767 , + -0.2642752 , -0.1780093 , -0.08995035, -0.5045371 , -0.31991875, + -0.3657458 , -0.51370513, -0.5282006 , -0.43932983, -0.44844607, + -0.4636747 , -0.541842 , -0.28016394, -0.5640383 , -0.28513128, + -0.44549647, -0.43519762, -0.42721725, -0.57123166, -0.40650263, + -0.12434424, -0.2684522 , -0.13034356, -0.41602165, -0.2656361 , + -0.63166195, -0.24252081, -0.49084774, -0.66598415, -0.21980134, + -0.20069678, -0.3857888 , -0.5750186 , -0.28702593, -0.42044863, + -0.43332833, -0.05599573, -0.17978549, -0.33969852, -0.42026928, + -0.19741155, -0.2731103 , -0.5647213 , -0.61824954, -0.44372314, + -0.16089505, -0.6141717 , -0.27717924, -0.39058053, -0.30320698, + -0.3017265 , -0.36256546, -0.29453766, -0.7215134 , -0.36851448, + -0.45460737, -0.25023326, -0.24756469, 0.5198105 , -0.49648827, + -0.7883031 , -0.40549907, -0.44536465, -0.34158024, -0.4282217 , + -0.45166638, -0.285461 , -0.5351079 , -0.5246095 , -0.38692436, + -0.4037612 , -0.68792176, 0.09384751, -0.3728323 , -0.54188645, + -0.14435874, -0.13926642, -0.48824954, 0.103251 , -0.5347266 , + -0.12768686, -0.2596208 , -0.33724293, 0.41711676, -0.46974605, + -0.31142974, -0.46029773, -0.38704032, -0.4837437 , -0.4762127 , + -0.48931426, -0.27646464, -0.7709319 , -0.11419559, -0.3580196 , + -0.41691482, -0.523539 , -0.7911689 , -0.14548191, -0.32703182, + -0.22812605, -0.5490665 , -0.27162734, -0.47042787, -0.27687344, + -0.29244688, -0.14437722, -0.42082107, -0.46603575, -0.3757466 , + -0.43942484, -0.35704526, -0.50712955, -0.06310245, -0.51886773, + -0.36505163, -0.29938784, -0.47636163, -0.798845 , -0.39453545, + -0.42888972, -0.34777957, -0.42616373, -0.37013018, -0.2707561 , + -0.22923009, -0.4484783 , 0.0781149 , -0.54366326, -0.31296468, + -0.46913856, -0.26365936, -0.05434812, -0.43664017, -0.262903 , + -0.21987678, -0.29727325, -0.53905886, -0.2881995 , -0.2697021 , + -0.56999284, -0.23401402, -0.25733256, -0.48550016, -0.53915155, + -0.29186913, -0.41333848, -0.11731004, -0.1876421 , -0.38366327, + -0.69185966, -0.54369533, -0.13784738, -0.45933396, -0.47800148, + -0.42821413, -0.37634867, -0.42712605, -0.06954894, -0.445592 , + -0.40655133, -0.3772172 , -0.7594351 , -0.5565112 , -0.25855637, + -0.13048875, -0.3732955 , -0.30561906, -0.5940574 , -0.41309616, + -0.314704 , -0.3010439 , -0.5837942 , 0.40323326, -0.24861278, + -0.4630232 , -0.538137 , -0.37435848, -0.46322653, -0.48135036, + -0.26899594, -0.17752594, -0.46429896, -0.62554 , -0.4089612 , + -0.34856302, -0.07396294, -0.4918398 , -0.44181666, -0.21933144, + -0.41612074, -0.46780694, -0.25013933, -0.3507257 , -0.31338322, + -0.49207342, -0.573551 , -0.25965518, -0.27235296, -0.46362886, + -0.572643 , -0.36132613, -0.5891415 , -0.43028024, -0.58089757, + -0.5628406 , -0.23266539, -0.48305756, -0.4171379 , -0.30668163, + -0.800837 , 0.16001745, -0.32928717, -0.7041144 , -0.28114614, + -0.1339794 , -0.1539137 , -0.48452973, -0.2916893 , -0.31748015, + -0.41399214, -0.5892567 , -0.6904999 , -0.39504105, -0.05394559, + -0.4220559 , -0.08093858, -0.55065817, -0.42636526, -0.36694276, + -0.6491974 , -0.16998939, -0.6064266 , -0.49446252, -0.25082645, + -0.3862861 , -0.5665641 , -0.12244188, -0.59443516, -0.43923312, + -0.03984782, -0.4539737 , -0.2804616 , -0.28886226, -0.41679534, + -0.7231176 , -0.44628784, -0.17072082, -0.55711204, -0.33631635, + -0.3513167 , -0.44696072, -0.09838888, -0.5331628 , -0.3244455 , + -0.22034623, -0.39987078, -0.52567166, -0.18180668, -0.1563768 , + -0.15176131, -0.34147656, -0.39736265, -0.22072102, -0.16304936, + -0.4934699 , -0.35189024, -0.38949656, 0.13233016, -0.2961783 , + -0.03038953, -0.19008778, -0.43870625, -0.14041252, -0.46124077, + -0.6101991 , -0.78746367, -0.41358054, -0.0965796 , -0.17454897, + -0.46205354, -0.5193903 , -0.43612164, -0.338933 , -0.2509465 , + -0.25011972, -0.32999432, -0.5068436 , -0.40737084, -0.59392023, + -0.08752135, -0.6887468 , -0.02706476, -0.34656492, -0.45208535, + 0.12221402, -0.07045163, -0.6397726 , -0.32004473, -0.2506177 , + -0.41133684, -0.20769219, -0.3077098 , -0.48824766, -0.59370387, + -0.3106976 , -0.09882683, -0.23733851, -0.5397758 , -0.384874 , + -0.48532894, -0.35305282, -0.44897798, -0.57191414, -0.23201816, + -0.40165707, -0.45165244, -0.30947968, -0.57236695, -0.40093747, + -0.610941 , -0.4310733 , 0.00399699, -0.17124976, -0.33956885, + -0.59238917, -0.4834356 , -0.23509572, -0.3706481 , -0.46995562, + -0.30272713, -0.42109653, -0.16173697, -0.3515516 , -0.46951428, + -0.24467751, -0.37127313], dtype=float32), + kernel: Array([[-0.16517928, -0.06662413, -0.32296306, ..., 0.17395326, + -0.00140762, -0.21535204], + [-0.2497043 , -0.32800463, -0.54542446, ..., -0.37079784, + -0.4540639 , -0.2762963 ], + [ 0.20276228, -0.2351077 , -0.4339836 , ..., -0.07039763, + 0.09680251, -0.00943627], + ..., + [-0.0865024 , -0.07723805, -1.0022317 , ..., -0.22161727, + 0.2631157 , -0.2520337 ], + [-0.06232752, 0.41687462, 0.43454677, ..., 0.17987384, + -0.22243185, 0.2732002 ], + [-0.1515276 , 0.08726911, -0.16446725, ..., -0.1428107 , + -0.12378144, -0.24571244]], dtype=float32), + }, + output: { + bias: Array([-2.27567896e-01, 3.73193085e-01, 1.41862273e-01, 8.60335305e-02, + 6.99756369e-02, -1.08260542e-01, -7.94294663e-03, 1.12044938e-01, + -1.97965074e-02, 1.03512242e-01, 1.25780692e-02, 2.14173347e-02, + -1.95132956e-01, -1.66688226e-02, -5.57698980e-02, 2.21790038e-02, + 8.23988672e-03, -1.56303987e-01, 2.77717095e-02, -6.28630817e-02, + -1.10570341e-01, 1.39788479e-01, -9.08712447e-02, -7.11519644e-03, + -1.21463031e-01, 1.24595746e-01, 6.99105812e-03, 6.04419075e-02, + 1.16523303e-01, -3.08515457e-03, 1.71046168e-01, -4.62258421e-02, + 1.37897981e-02, -9.23820660e-02, 1.22136414e-01, -5.12558557e-02, + -4.16860264e-03, 5.67508861e-02, -8.22613947e-04, -1.72597811e-01, + -1.73500534e-02, -1.02644004e-01, -7.16797352e-01, 6.50356859e-02, + -9.02397409e-02, -7.97297060e-02, -3.20093259e-02, 7.92354718e-02, + -5.91707882e-03, -5.96906338e-03, 2.83249887e-03, 4.21253145e-02, + -1.67518064e-01, -2.99427286e-02, -5.09011634e-02, -3.20116989e-02, + -1.33028731e-01, 2.27317363e-02, 1.29082099e-01, 1.18409574e-01, + 1.57172829e-02, -2.19231278e-01, -2.51264381e-03, -4.79261391e-02, + 4.22019213e-02, -7.78303593e-02, -7.12049752e-02, -8.38479102e-02, + -3.88463922e-02, -3.15606929e-02, -6.03214689e-02, 4.97607961e-02, + -3.79606709e-02, 7.08331587e-04, 9.65211987e-01, 1.56220675e-01, + -2.14344375e-02, -2.15863492e-02, -1.17833465e-01, 1.19182810e-01, + -2.71626860e-01, 1.97731983e-03, 3.98320071e-02, 2.02119559e-01, + 6.09681681e-02, -3.32502718e-03, 1.39364108e-01, 5.88153526e-02, + -4.45023738e-02, -6.55459762e-02, 4.42181528e-02, -8.78035650e-02, + -4.80993986e-02, -1.09087467e-01, -9.66937765e-02, -1.30174085e-01, + 2.12008785e-03, 2.63354164e-02, -9.37493294e-02, 7.85378218e-02, + 1.73281968e-01, -7.77669847e-02, 1.33560196e-01, -1.09741293e-01, + 1.85674503e-01, 1.03994340e-01, -1.89410765e-02, -1.18796170e-01, + -1.62400916e-01, 2.00676285e-02, -1.06416032e-01, 1.79987997e-02, + 1.19247243e-01, -2.42678579e-02, 1.64977908e-01, -7.41590792e-03, + -1.76154330e-01, 2.93482076e-02, -6.35338202e-02, 2.07892321e-02, + 8.92439187e-02, 2.59270482e-02, 2.71555986e-02, -7.91539550e-02, + 9.35807526e-02, 5.84603362e-02, 1.93760265e-02, 1.72683045e-01], dtype=float32), + kernel: Array([[ 6.6307461e-01, -2.6014799e-01, -8.3183780e-02, ..., + -2.6699215e-01, -1.0033758e-01, -2.1524315e-01], + [-1.4059058e-01, -4.3626988e-01, -8.1478050e-03, ..., + -3.1092042e-01, 3.1286913e-01, -3.7179798e-02], + [ 7.0444113e-01, -1.0556934e-05, -2.0829076e-02, ..., + -9.0309918e-01, 3.2471472e-01, 5.0782524e-02], + ..., + [-7.5270250e-02, 3.7796192e-02, -5.1213449e-01, ..., + 3.4989411e-01, -1.0396782e-01, -5.7098109e-01], + [ 9.8383605e-02, -1.5289782e-01, 3.3700535e-01, ..., + 5.9063655e-01, -9.2052268e-03, -1.1524218e-01], + [ 1.1272128e-01, -6.6285774e-02, -2.5734550e-01, ..., + -1.2502365e-01, 9.1840588e-02, 2.3722002e-01]], dtype=float32), + }, + }, + }, + }, + 1: { + attention: { + attention: { + key: { + bias: Array([ 0.08042084, 0.17798816, 0.05449522, 0.09923156, 0.02864883, + -0.06706466, 0.12987332, 0.05871709, 0.05153077, 0.05436414, + 0.05116336, -0.14068684, -0.11056656, -0.00991776, 0.00717054, + 0.0770971 , 0.11563851, 0.02029003, 0.02723394, -0.12056217, + -0.06265891, -0.05970899, 0.14154014, -0.03163632, 0.0481686 , + 0.02711394, 0.04412205, 0.02307438, -0.05139432, 0.03842278, + -0.07357287, -0.08965195, -0.08623472, 0.03718801, 0.84410733, + 0.4109549 , 0.3907228 , -0.44974568, -0.3517945 , 0.03430397, + -0.07853877, -0.30442524, 0.5081588 , 0.15563154, 0.17845473, + 0.14964798, 0.01935043, 0.78861445, 0.13657925, 0.9571119 , + -0.15077294, 0.04331086, -0.23139133, -0.09555604, 0.06358018, + -0.59850585, -0.37260893, -0.9738547 , -0.01462804, 0.12405673, + 0.07412421, -0.01379338, 0.05929387, 0.7805736 , -0.36394337, + 0.47818112, -0.43866792, 0.63309664, 0.1147754 , 0.45001298, + 0.3213029 , 0.44093522, 0.51016414, -0.18746474, -0.13791506, + 0.04778931, -0.49275073, 0.04839996, -0.38444966, -0.4882343 , + -0.6615728 , 0.09630081, -0.29930323, -0.34863642, 0.77798736, + -0.55292934, -0.56388927, 0.09340773, -0.4621951 , 0.26494777, + 0.31161246, 0.44366807, 0.6213311 , 0.01337955, 0.14917652, + -0.4876238 , 0.06336308, -0.12776731, -0.07288839, -0.18272795, + 0.23383974, -0.059283 , -0.08504926, -0.12375498, 0.3620917 , + 0.04436552, 0.15867496, 0.12362888, -0.09505215, 0.18897517, + 0.08316641, -0.18332016, -0.26461446, 0.17972536, 0.23397446, + -0.04752706, -0.04319599, 0.12071403, 0.25849587, 0.06566656, + 0.19910227, 0.22742352, -0.22645345, -0.14635295, -0.2507506 , + -0.18631294, -0.05247986, 0.12347726], dtype=float32), + kernel: Array([[-0.06023952, 0.297179 , -0.02699618, ..., -0.1145321 , + -0.08048212, -0.10587901], + [ 0.6934901 , 0.47458088, 0.61580837, ..., -0.13140583, + -0.14733373, 0.26471645], + [ 0.09990517, -0.17357504, 0.05461919, ..., -0.11821933, + -0.24399522, -0.15253429], + ..., + [-0.2205276 , -0.23215702, -0.24366312, ..., -0.03175567, + 0.22550093, 0.27891138], + [-0.24111255, -0.27693352, -0.21988523, ..., 0.30499363, + 0.37258562, -0.13957877], + [-0.103215 , -0.14666593, -0.04492585, ..., -0.6473914 , + -0.23865867, 0.19981082]], dtype=float32), + }, + query: { + bias: Array([-0.11081476, -0.09144565, -0.06050301, -0.03483055, 0.03535065, + -0.1849559 , -0.04721399, -0.07624152, -0.10357637, 0.22812082, + -0.3447222 , 0.0331145 , 0.05764522, -0.13265188, 0.1981749 , + -0.03139308, -0.07606025, 0.13599226, 0.02934942, 0.12725802, + 0.07556148, 0.07594807, -0.04966472, -0.15265177, -0.35979733, + 0.01576409, -0.07025045, -0.19974303, -0.22990641, 0.0952646 , + -0.33106107, 0.04587565, -0.41658264, 0.33157024, -0.51854265, + 0.34642565, -0.38225 , 0.21841964, 0.4542027 , -0.3708585 , + -0.17867757, 0.38574484, -0.40515402, -0.18270488, 0.15764216, + -0.09051255, -0.47678202, -0.5485276 , 0.17880292, -0.42515096, + 0.11444188, 0.27703342, -0.23375264, -0.17498691, 0.13194664, + 0.06895838, 0.32982063, 0.17495483, -0.16514482, 0.16266893, + 0.24224894, 0.49473688, -0.7105683 , -0.22717172, -0.4628581 , + -0.02517562, 0.13951837, 0.04581514, -0.02814059, 0.06052845, + 0.89534795, -0.01086807, -0.2937633 , 0.4263962 , 0.01343843, + 0.9783011 , 0.23017581, 0.7163039 , 0.41632223, 0.17262603, + -0.15068051, 0.850449 , -0.16111681, -0.02732765, -0.18101065, + 0.47148392, 0.18522678, -0.06222672, 0.15925601, -0.07961265, + 0.71117395, 0.02433823, -0.35966122, -0.09352748, -0.06117439, + -0.13183966, -0.08041007, -0.27098748, -0.71445733, -0.5658294 , + 0.51616377, -0.56782836, -0.68260497, -0.5732899 , 0.5585706 , + -0.14336808, 0.59159863, -0.21151459, 0.42952493, 0.46055183, + 0.46279824, -0.4195606 , -0.5513547 , 0.34171253, 0.69069153, + -0.25621027, -0.46899697, 0.588077 , 0.6255269 , 0.5819053 , + 0.54691887, 0.59542644, -0.5749214 , -0.57993734, 0.42467687, + -0.70540833, -0.34917867, 0.3106283 ], dtype=float32), + kernel: Array([[-0.11731137, -0.1564569 , -0.11458346, ..., 0.02826559, + -0.03335454, 0.14976554], + [-0.04609706, -0.04719561, -0.02786452, ..., -0.30554634, + -0.20783764, -0.11992735], + [ 0.0489803 , 0.12078558, 0.0699265 , ..., 0.26866522, + -0.40512353, -0.5030211 ], + ..., + [-0.00423968, -0.00703923, -0.03773196, ..., -0.10513086, + -0.145968 , 0.07945077], + [ 0.07510082, -0.13710845, 0.03422854, ..., 0.15533802, + 0.8349194 , -0.10965192], + [-0.07718561, 0.09582806, -0.06439573, ..., -0.12381844, + -0.9082412 , 0.04665521]], dtype=float32), + }, + value: { + bias: Array([-0.11375186, 0.01737056, -0.11132707, -0.01888963, 0.02005741, + -0.00092973, -0.11860004, 0.14434907, -0.06270761, 0.14292052, + 0.14930397, 0.06040736, 0.1552559 , -0.06054825, 0.13105005, + 0.05230521, 0.03080181, -0.14906318, 0.04209505, -0.11033 , + -0.02365249, -0.14162254, -0.00175851, 0.03008844, 0.04629046, + 0.10597341, 0.02212737, 0.13431458, -0.03534039, 0.00220095, + -0.05828169, 0.01887055, 0.07417873, -0.18334955, 0.05120145, + 0.03194632, 0.09620124, 0.07511445, 0.14970838, -0.3099061 , + -0.00785524, 0.21368372, -0.24593364, -0.20003676, 0.29439747, + 0.02457003, -0.0362221 , -0.04669455, -0.11111037, -0.09324432, + -0.007548 , -0.05471316, -0.1160496 , -0.08286496, 0.10508952, + -0.29490575, 0.00155351, 0.13688038, 0.055782 , -0.241635 , + 0.19978897, 0.12800941, -0.08591252, -0.17519936, 0.0762416 , + 0.03798862, -0.04639675, -0.1959964 , -0.25277418, 0.169555 , + -0.06163629, -0.11278319, -0.01159456, -0.06765992, 0.0005254 , + -0.06396206, -0.00113279, 0.1249745 , -0.02429713, -0.01882967, + -0.04579956, -0.02704 , 0.00707008, 0.02282375, -0.10874218, + -0.01254413, 0.05641572, -0.06401654, 0.17181516, 0.0763936 , + -0.07932342, -0.12056571, 0.0820564 , -0.01728764, 0.04256084, + -0.05678854, 0.11785327, 0.2377933 , -0.19259618, -0.08974701, + -0.19949402, -0.09832341, 0.12747605, 0.01983653, 0.00941677, + 0.11338363, -0.20597233, -0.04107258, 0.02300076, -0.04743035, + 0.11633945, 0.12718917, 0.23441505, 0.01317911, -0.19377421, + -0.06671913, 0.04473834, 0.00467375, -0.15131958, -0.07110944, + -0.16257668, -0.08920506, 0.03026637, -0.0935519 , 0.06253995, + -0.19525133, -0.21808124, -0.19407861], dtype=float32), + kernel: Array([[-0.01158661, 0.04303771, -0.05590956, ..., -0.05333043, + -0.00572822, -0.01724763], + [ 0.2719801 , 0.14965016, -0.02898863, ..., -0.06263707, + -0.35550225, -0.28345698], + [-0.05632893, -0.26717782, -0.07165924, ..., -0.43430653, + 0.0922191 , 0.10500429], + ..., + [-0.04451123, -0.00324844, 0.15866533, ..., -0.15973136, + -0.06277192, -0.10269681], + [-0.28674033, 0.40732047, -0.23387024, ..., 0.5245572 , + -0.20356956, -0.19302332], + [-0.02055671, -0.35077375, 0.13905095, ..., 0.08351879, + 0.4497927 , 0.44072792]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-2.29397908e-01, 3.42801064e-01, 1.12294301e-01, -1.40454136e-02, + 3.03802863e-02, -1.64602742e-01, 6.72679991e-02, 9.77464989e-02, + -2.96701267e-02, 7.89122283e-02, 5.59587255e-02, 1.48079200e-02, + -1.85802028e-01, -3.27114910e-02, -2.40307394e-02, 6.82050884e-02, + 5.04146097e-04, -1.56787038e-01, 6.75788894e-02, -1.31789781e-02, + -4.55502272e-02, 1.27612054e-01, -7.01643378e-02, -2.48179883e-02, + -1.31983072e-01, 1.38503700e-01, 1.17344763e-02, 7.56531805e-02, + 6.57847598e-02, 2.69352719e-02, 2.21325293e-01, -5.05222380e-02, + 3.61935459e-02, -1.04897447e-01, 5.24883904e-02, -8.26861486e-02, + 1.44010643e-02, 8.25399533e-02, -6.71743648e-03, -1.94458306e-01, + -2.49852352e-02, -8.23389366e-02, -6.81263685e-01, 8.11951086e-02, + -8.49441960e-02, -1.40838427e-02, -1.44380983e-02, 6.43714517e-02, + -1.57576092e-02, -3.63134891e-02, -1.25548737e-02, 1.41190467e-02, + -1.39931515e-01, 1.46125898e-01, -4.10604011e-03, -3.28917056e-02, + -1.47544041e-01, -7.30510205e-02, 1.31265923e-01, 3.08194682e-02, + 5.00784516e-02, -2.18488395e-01, 2.87581724e-03, -3.31245400e-02, + 2.70526409e-02, -1.38555646e-01, -1.16920797e-02, -7.85453171e-02, + 1.44379027e-02, -8.10271129e-02, -7.89678767e-02, 6.87262276e-03, + 2.95094848e-02, -2.96989549e-02, 8.52790475e-01, 1.68656245e-01, + -4.95392680e-02, 1.01274252e-03, -1.17389508e-01, 1.42783284e-01, + -2.95234203e-01, -3.66869681e-02, 2.10161004e-02, 1.61923394e-01, + 1.00733116e-01, 9.56858695e-02, 1.62279665e-01, 6.97396919e-02, + -6.76168874e-02, -1.21831730e-01, 2.96244565e-02, -8.43596309e-02, + -4.85159941e-02, -1.54231772e-01, -8.46877992e-02, -9.01621133e-02, + -7.94285350e-03, 9.12197307e-03, -5.20062111e-02, 2.70500332e-02, + 1.75196305e-01, -7.51989707e-02, 1.59212708e-01, -8.15638751e-02, + 1.15850195e-01, -1.62196823e-03, 6.47357618e-03, -1.16396412e-01, + -1.10036433e-01, -1.06435921e-02, -1.69102713e-01, 2.04750933e-02, + 1.00554541e-01, -3.42080109e-02, 1.30484760e-01, 1.22546069e-02, + -1.68031573e-01, 3.20073776e-02, -6.57493770e-02, -7.96901435e-03, + 1.29959419e-01, 8.22199881e-02, 4.45892625e-02, -4.60687764e-02, + 6.53466284e-02, 2.84257364e-02, -7.09084561e-03, 8.53643045e-02], dtype=float32), + kernel: Array([[-4.33085859e-01, -5.61205745e-01, 2.93745458e-01, ..., + -1.06319696e-01, -3.59772116e-01, -5.50748073e-02], + [ 1.29990995e-01, -1.94321834e-02, 3.05669427e-01, ..., + 1.81368724e-01, 9.14699510e-02, -6.21524341e-02], + [-2.22445935e-01, 1.21082552e-01, -5.40154229e-04, ..., + -1.46840677e-01, -1.43376559e-01, 9.97718945e-02], + ..., + [ 6.94694817e-01, 8.00504535e-03, 7.26794839e-01, ..., + -2.44273514e-01, -3.41350943e-01, -1.44506514e-01], + [ 6.64600372e-01, -1.74710259e-01, 2.86397934e-01, ..., + -4.65431027e-02, -2.18760711e-03, -2.81117141e-01], + [ 5.85957587e-01, -1.68455601e-01, 3.16843003e-01, ..., + -4.89029028e-02, -3.35241156e-03, -2.94326484e-01]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([ 6.88922033e-02, 5.80978803e-02, 3.85873437e-01, -3.11470646e-02, + 1.69555983e-03, -2.32034996e-01, -1.58923380e-02, -3.62700596e-02, + 5.75650949e-04, -3.82705927e-01, 2.31035739e-01, 6.68624789e-02, + -1.72114998e-01, -5.66142239e-02, 2.09674109e-02, -2.80936039e-03, + -2.27225497e-02, -3.59404653e-01, 1.31325588e-01, -3.94533155e-03, + -1.76445603e-01, 4.06457812e-01, -3.71316727e-03, -2.46738438e-02, + -1.81620270e-01, 4.64766249e-02, 7.05432147e-02, 7.28050619e-02, + -5.19782789e-02, -4.45098169e-02, 3.48823667e-01, 6.42984435e-02, + -3.04025542e-02, -9.68929306e-02, -5.14351986e-02, -3.97855928e-03, + -3.20420787e-02, 2.67196506e-01, -7.22298101e-02, -2.07284018e-01, + 4.90039028e-02, -2.19781891e-01, -1.35466352e-01, 8.54400694e-02, + 1.19789563e-01, -1.32402414e-02, -2.06071258e-01, 3.45563710e-01, + 1.00530580e-01, 8.49714726e-02, 3.91122419e-03, 2.15060711e-02, + 8.78436714e-02, 2.07931370e-01, 3.97193991e-02, -1.13610223e-01, + 3.73787954e-02, 1.23064490e-02, -9.43080038e-02, -1.86623499e-01, + -3.85700941e-01, -1.63304899e-02, 1.10437632e-01, -1.03100136e-01, + -6.19614311e-03, -3.90111618e-02, 1.53294563e-01, -1.00014314e-01, + 8.98484327e-03, -1.33218288e-01, 1.19500153e-01, -1.21143341e-01, + 2.10582778e-01, 1.27065793e-01, -1.65820807e-01, 2.54214227e-01, + 5.12329042e-02, -9.36500803e-02, 5.69578595e-02, 3.65398824e-01, + -1.66889839e-02, -8.60542282e-02, -2.17619090e-04, 7.38569945e-02, + 1.74466029e-01, 7.53093883e-02, -1.50982007e-01, -1.61455601e-01, + -4.24303383e-01, -1.41928092e-01, 1.67686343e-02, -2.84632683e-01, + 1.45200789e-01, 5.28667821e-03, -2.16334872e-02, 4.74276505e-02, + 1.93778221e-02, -1.02245426e-02, 8.27698484e-02, -4.25329745e-01, + -9.56311524e-02, 2.45452020e-02, 5.58077872e-01, 8.93760026e-02, + -3.42835188e-02, -7.27823749e-03, -4.62576747e-02, -1.11996412e-01, + -2.86399245e-01, 3.93988378e-03, -5.40265851e-02, 1.06320893e-02, + -1.11137971e-01, 9.71371382e-02, 7.98522159e-02, -7.13857710e-02, + -4.83553410e-01, 1.70313239e-01, -1.70816466e-01, 8.59201849e-02, + -3.64008136e-02, -2.24960689e-02, -3.12829800e-02, 1.58451021e-01, + 9.86519852e-04, -3.24616969e-01, -2.68537730e-01, -1.29793763e-01], dtype=float32), + scale: Array([ 0.16587867, 0.62729955, 0.78539467, 0.33312872, 0.02448054, + 0.52365166, 0.09352965, 0.57490396, -0.00181935, 0.62166315, + 0.5622861 , 0.632528 , 0.32434413, 0.33608988, 0.09364542, + 0.00295282, 0.20884101, 0.4969771 , 0.49006283, 0.3390643 , + 0.76553243, 0.75735354, 0.64736444, 0.20501752, 0.5099314 , + 0.39772737, 0.4760952 , 0.2673762 , 0.5248422 , 0.34622046, + 0.67513955, 0.12790139, 0.78136003, 0.46376252, 0.17142247, + 0.02035132, 0.35290378, 0.79245657, 0.43180364, 0.58383363, + 0.28102416, 0.6248626 , 0.09244277, 0.38200957, 0.36086422, + 0.44281945, 0.62160456, 0.49688524, 0.47265998, 0.4368852 , + 0.00439632, 0.24748251, 0.2077039 , 0.71143705, 0.55966383, + 0.76935214, 0.2611507 , 0.3572642 , 0.4780699 , 0.5534745 , + 0.5650403 , 0.49476576, 0.6574258 , 0.38081974, 0.15224549, + 0.2871514 , 0.5269491 , 0.5443608 , 0.47709846, 0.6731636 , + 0.48256594, 0.57757723, 0.52946705, 0.27599993, 0.11646109, + 0.7459179 , 0.6025415 , 0.4426929 , 0.150529 , 0.6974664 , + 0.07860193, 0.63189274, 0.6666647 , 0.8469711 , 0.4019909 , + 0.45879367, 0.22790541, 0.5310602 , 0.46144342, 0.46353984, + 0.3125676 , 0.6395497 , 0.5902147 , -0.00431166, 0.51785195, + 0.44832194, 0.2982607 , -0.02292351, 0.48743743, 0.42023692, + 0.5205952 , 0.08753084, 0.76318944, 0.42676944, 0.6018406 , + 0.7163417 , 0.55152804, 0.8593687 , 0.4082306 , 0.4901796 , + 0.64518034, 0.5525667 , 0.34930795, 0.46910506, 0.6457558 , + 0.54612505, 0.61357117, 0.49350467, 0.41991618, 0.55927193, + 0.4409146 , 0.15714881, 0.31889114, 0.26255667, 0.45123503, + 0.56282735, 0.50075245, 0.34856102], dtype=float32), + }, + layernorm_before: { + bias: Array([ 2.82420870e-02, -2.29402691e-01, -2.96053607e-02, -1.30823404e-02, + -2.65389867e-03, 2.99529568e-03, -3.38511716e-04, -4.37224060e-02, + 2.60172427e-01, -2.05475893e-02, 1.60881188e-02, -1.69427078e-02, + -5.21531154e-04, 7.41502224e-03, -1.61934067e-02, 3.54598239e-02, + -1.85820572e-02, 2.71367803e-02, -4.88283485e-03, -6.32894114e-02, + 3.68127339e-02, -6.04641587e-02, 1.19915558e-03, 1.77571445e-03, + -7.46919066e-02, -1.53914511e-01, -3.68062826e-03, 1.60319414e-02, + 6.28713518e-02, 3.06911282e-02, -2.31577735e-02, -6.57155272e-03, + 6.15396686e-02, 2.63886862e-02, -4.85325791e-02, 7.88616296e-03, + 8.01346614e-04, -1.55313745e-01, -2.11913697e-03, -6.54404704e-03, + 5.55190295e-02, -7.31622195e-03, 1.96012199e-01, 5.94329368e-03, + -2.78697144e-02, -4.93782759e-02, 3.01586390e-02, 6.05354905e-02, + -1.21998333e-03, 3.04064993e-02, 1.04687298e-02, 1.09375557e-02, + -2.71287002e-03, -2.32986733e-01, 1.44290384e-02, -2.30109710e-02, + 4.33547981e-02, 5.82204685e-02, -2.50215475e-02, 1.38790637e-01, + -6.68904185e-02, 1.24691881e-01, -4.73514060e-03, -3.82121443e-03, + 3.83599102e-02, 8.80813003e-02, -1.65379103e-02, -7.19220552e-04, + 3.60673331e-02, 1.00804691e-03, -1.36750117e-02, -2.13624742e-02, + 3.65178101e-02, 2.25900616e-02, -4.44582880e-01, -2.65965145e-02, + -5.42465337e-02, -7.62177110e-02, 6.57133060e-04, -1.34424558e-02, + 3.63113172e-02, 7.55314380e-02, 8.09766643e-04, -2.51245238e-02, + -2.74020876e-03, -4.82323244e-02, -7.02568442e-02, 1.68840366e-03, + 1.02340458e-02, 1.46568000e-01, -2.30909362e-02, 4.56218012e-02, + -6.14786409e-02, -2.31754649e-02, 1.12327132e-02, 8.98616686e-02, + 3.33565250e-02, 3.43459705e-03, -1.15024485e-02, -9.24066082e-02, + -3.95509861e-02, -2.51406804e-03, -9.44837704e-02, -6.91622263e-04, + -3.78696583e-02, -7.13908523e-02, -6.70393333e-02, 2.14871876e-02, + -3.32830548e-02, -9.49056670e-02, 3.31703201e-02, 5.51439868e-03, + -1.63990129e-02, 1.68498412e-01, -1.52823254e-01, 1.28563598e-01, + 2.69770622e-02, -2.52538477e-03, -6.20844103e-02, -2.31217127e-03, + -1.12146828e-02, 2.15321942e-03, -1.07783070e-02, 7.04224780e-03, + -5.70042757e-04, -1.27387140e-02, -1.50637852e-03, -1.26768842e-01], dtype=float32), + scale: Array([ 0.02225583, 0.3582427 , 0.18810613, 0.2897796 , 0.01008445, + 0.20093079, 0.00262845, 0.18910973, 0.20207141, 0.05451531, + 0.13650979, 0.22285123, 0.04530312, 0.00609608, 0.10985828, + 0.09655713, 0.13278736, 0.08787937, 0.02501552, 0.16963154, + 0.22782324, 0.20419084, 0.23176458, 0.04793768, 0.12091924, + 0.15503103, 0.12642743, 0.07698647, 0.37561363, 0.1734934 , + 0.22910334, 0.00436596, 0.25682634, 0.24328624, 0.2714059 , + 0.06567708, 0.02274263, 0.567844 , -0.00143492, 0.16284508, + 0.18691777, 0.3501844 , 0.11066452, 0.01387185, 0.17211044, + 0.31377098, 0.1559427 , 0.31400132, 0.03055096, 0.1480449 , + 0.06916483, 0.05761943, -0.00319004, 0.38499779, 0.0325634 , + 0.3114719 , 0.07862399, 0.19379751, 0.12478653, 0.41179135, + 0.26195255, 0.3847175 , 0.16452418, 0.00272222, 0.15891619, + 0.11142635, 0.14923933, 0.10400768, 0.21697171, 0.12220434, + 0.24840005, 0.16391885, 0.09268359, 0.12174353, 0.12482438, + 0.2716236 , 0.22829536, 0.2231431 , -0.0092565 , 0.03284183, + 0.06326266, 0.15956563, 0.06667496, 0.22676413, 0.01954163, + 0.2572471 , 0.14380816, 0.3025213 , 0.13769881, 0.38677344, + 0.05833767, 0.11660744, 0.1784911 , 0.31127843, 0.0378647 , + 0.11842007, 0.13203348, 0.08775622, 0.22929348, 0.1685133 , + 0.17730162, 0.14410907, 0.3452061 , -0.01377543, 0.44221473, + 0.43522978, 0.29845122, 0.13120209, 0.45809656, 0.19538262, + 0.46633697, 0.1401029 , 0.05973453, 0.4215374 , 0.16866335, + 0.27390364, 0.03285838, 0.01355824, 0.18128489, 0.2254643 , + 0.03559769, 0.00747606, 0.18407959, 0.02931737, -0.00183865, + 0.0155194 , 0.2910443 , 0.2715314 ], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-2.87663639e-01, -2.04487324e-01, -3.44864547e-01, -2.00529769e-01, + -5.94829321e-01, -1.03927039e-01, -2.77930766e-01, -4.94999826e-01, + -2.08718747e-01, -3.45035255e-01, -2.67087758e-01, -4.77189213e-01, + -1.66359320e-01, -4.48330075e-01, -4.78699297e-01, -2.21549183e-01, + -2.41500542e-01, -3.03187490e-01, -4.03688133e-01, -3.62672597e-01, + 3.11330687e-02, -4.75372255e-01, -1.33399740e-01, -1.58057332e-01, + -1.38183832e-01, -3.53402674e-01, -4.83101279e-01, -3.49856049e-01, + -4.62018728e-01, -2.20034316e-01, -1.45977482e-01, -1.44924179e-01, + -3.43084693e-01, -2.04671875e-01, -2.16533989e-01, -5.23977339e-01, + -2.12659031e-01, -2.74191111e-01, -2.65450865e-01, -4.38724220e-01, + -4.05285835e-01, -4.25657749e-01, -2.83261627e-01, -1.07905425e-01, + -3.10509741e-01, -3.69468153e-01, -5.25859118e-01, -3.74539018e-01, + -4.44122553e-01, -2.76227951e-01, -3.62340897e-01, -3.20163608e-01, + -5.22766829e-01, -3.68136406e-01, -6.11795008e-01, -1.88745782e-01, + -4.54058141e-01, -1.91686556e-01, -4.25439298e-01, -5.22909224e-01, + -3.05993110e-01, -3.57376486e-01, -3.10534716e-01, -2.63750911e-01, + -5.45675814e-01, -4.23284829e-01, -1.21132560e-01, -3.99448633e-01, + -1.06209874e-01, -3.28158796e-01, -5.71545482e-01, -1.91518441e-01, + -5.93747199e-01, -3.57770473e-01, -2.59463459e-01, -3.54470700e-01, + -5.28468788e-01, -3.18399936e-01, -4.58652407e-01, -2.15082556e-01, + -4.01729345e-01, -5.17883539e-01, -3.26544613e-01, -2.71224737e-01, + -3.98273647e-01, -3.58481169e-01, -4.33161110e-01, -4.17540729e-01, + -4.41764146e-01, -4.13900316e-01, -4.20035869e-01, -4.61230427e-01, + -5.54932415e-01, -2.82214522e-01, -4.63347554e-01, -1.69810832e-01, + -4.96465772e-01, -3.57517421e-01, -3.08119923e-01, -5.05688906e-01, + -4.42314088e-01, -2.96192914e-01, -3.34470212e-01, -1.58352673e-01, + -1.96208924e-01, -1.71411768e-01, -2.78193712e-01, -4.14353907e-01, + -5.30355275e-01, -6.04073517e-02, -3.42649519e-01, -3.25455546e-01, + -2.80562103e-01, -3.16360086e-01, -1.88339218e-01, -1.45719707e-01, + -2.42292494e-01, -4.38852787e-01, -3.38206112e-01, -3.96864712e-01, + -3.94322276e-01, -5.04216313e-01, -4.47647184e-01, -3.14450145e-01, + -4.03099447e-01, -3.31715465e-01, -5.11624277e-01, -2.97003418e-01, + -1.71126753e-01, -5.23082614e-01, -3.97670537e-01, -3.33513230e-01, + -2.73618609e-01, -2.16470793e-01, -3.16209495e-01, -3.54464471e-01, + -4.11175519e-01, -3.47123802e-01, -4.30214375e-01, -2.96357393e-01, + -3.50862324e-01, -4.91615891e-01, -2.99721777e-01, -3.46929997e-01, + -6.17211461e-01, -3.17567289e-01, -2.59017050e-01, -5.43685555e-01, + -3.07671517e-01, -3.49160552e-01, -5.83933473e-01, -2.99324632e-01, + -3.56959522e-01, -4.75278974e-01, -4.69556153e-01, -4.00068551e-01, + -3.57878387e-01, -2.61400610e-01, 5.89218475e-02, -4.75838035e-01, + -3.18645447e-01, -2.91967630e-01, -3.57156008e-01, -5.22750139e-01, + -1.99007168e-01, -4.69590366e-01, -3.96272272e-01, -2.64699906e-01, + -3.60923529e-01, -1.91638634e-01, -3.21910053e-01, -2.48550743e-01, + -3.98402393e-01, -4.65990812e-01, -3.43011886e-01, -3.76989871e-01, + -3.99559021e-01, -3.55907410e-01, -4.88500834e-01, -3.14257860e-01, + -5.45107484e-01, -1.72966227e-01, -3.69181097e-01, -2.42310494e-01, + -3.32116067e-01, -3.86646032e-01, -2.56473958e-01, -3.48980218e-01, + -3.98971230e-01, -2.64990151e-01, -4.63962048e-01, -1.67646125e-01, + -3.03776145e-01, -2.27759063e-01, -4.02701110e-01, -3.71156693e-01, + -4.49759185e-01, -4.83259469e-01, -4.44558293e-01, -3.49739552e-01, + -4.34326917e-01, -2.64975369e-01, -2.89425910e-01, -4.49308783e-01, + -2.49016702e-01, -4.31563139e-01, -4.08892900e-01, -3.69399726e-01, + -6.46372512e-02, -3.25122327e-01, -1.55182466e-01, -1.11164600e-01, + -4.79935288e-01, -5.89126289e-01, -5.05477190e-01, -3.91809165e-01, + -4.87700224e-01, -5.42157114e-01, -5.01625121e-01, -3.88096064e-01, + -4.09615874e-01, -3.30858119e-03, -3.16691220e-01, -4.25819814e-01, + -2.07181290e-01, -2.52645940e-01, -4.78967577e-01, -4.66646552e-01, + -5.08263230e-01, -2.38753840e-01, -3.27672154e-01, -5.24233878e-01, + -3.54793280e-01, -3.08880091e-01, -2.62713552e-01, -3.00624251e-01, + -4.21428025e-01, -4.42182183e-01, -5.11559665e-01, -2.42620170e-01, + -3.54224473e-01, -3.06683362e-01, -4.52630949e-04, -4.44990546e-01, + -3.27576101e-01, -2.59946615e-01, -5.11295557e-01, -6.64835870e-01, + -1.85168549e-01, -5.48847735e-01, -2.54403502e-01, -4.44106311e-01, + -2.09086508e-01, -5.46322465e-01, -4.56374437e-01, -2.26774529e-01, + -3.03409159e-01, -2.95624912e-01, -4.80216295e-01, -1.95244759e-01, + -2.66378999e-01, -3.54911506e-01, -3.88011873e-01, -2.92912245e-01, + -4.73809302e-01, -2.48538107e-01, -2.53374189e-01, -5.34175217e-01, + 1.01788407e-02, -4.37775195e-01, -3.74520659e-01, -4.16397095e-01, + -2.86437750e-01, -1.17193200e-01, -4.65466827e-01, -2.45054185e-01, + -1.55902386e-01, -2.60219306e-01, -3.81529987e-01, -4.36873585e-01, + -3.58277500e-01, -4.33071882e-01, -4.30247694e-01, -3.75352204e-01, + -3.81595612e-01, -4.86366630e-01, -4.01012272e-01, -2.53684759e-01, + -3.21343958e-01, -5.21693230e-01, -2.32155509e-02, -3.01609904e-01, + -2.36861959e-01, -3.18232119e-01, -1.21864028e-01, -3.91450346e-01, + -5.88354170e-01, -3.05314749e-01, -3.78017843e-01, -9.56703201e-02, + -3.85766625e-01, -2.07299024e-01, -8.47784728e-02, -3.33295763e-01, + -4.15781707e-01, -4.48703825e-01, -1.40794486e-01, -1.18951991e-01, + -2.74386436e-01, -4.96934414e-01, -3.64418119e-01, -3.70397896e-01, + -1.35903224e-01, -1.35731950e-01, -3.99609268e-01, -2.93618023e-01, + -3.04051191e-01, -5.70618927e-01, -5.02060175e-01, -2.81283766e-01, + -6.31259501e-01, -3.06242526e-01, -2.12949291e-01, -4.86519247e-01, + -1.59262896e-01, -4.03478205e-01, -1.81495771e-01, -4.62975919e-01, + -3.86971891e-01, -5.12782514e-01, -2.76577592e-01, -2.94954062e-01, + -2.17922330e-01, -4.69827503e-01, 1.06356489e-02, -3.15419227e-01, + -3.53538573e-01, -5.48389792e-01, 5.30533865e-02, -3.23666394e-01, + -3.68639082e-01, -3.19764733e-01, -4.70012695e-01, -9.86519307e-02, + -3.48372728e-01, -2.50590980e-01, -3.00890654e-01, -4.49694991e-01, + -3.42396975e-01, -2.40296051e-01, -1.83449000e-01, -4.94159639e-01, + -3.83832961e-01, -2.35955358e-01, -4.39356416e-01, -2.83128679e-01, + -5.06890595e-01, -4.84923810e-01, -3.48398417e-01, -8.36336315e-02, + -3.36799473e-01, -4.19064090e-02, -5.25771677e-01, -4.51774269e-01, + -4.36057091e-01, -4.80761558e-01, -3.56559277e-01, -4.32838291e-01, + -4.25217807e-01, -2.08370864e-01, -6.56363890e-02, -5.52178204e-01, + -4.09472764e-01, -1.91283375e-01, -1.69247791e-01, -3.70922804e-01, + -4.36929941e-01, -1.94023415e-01, -2.98483998e-01, -1.78429320e-01, + -2.56660700e-01, -3.49405706e-01, -2.60902107e-01, -1.27692908e-01, + -2.62808502e-01, -3.31191778e-01, -5.31970501e-01, -2.73841202e-01, + -3.46037507e-01, -2.07595572e-01, -3.45578700e-01, -3.25063765e-01, + -3.59042734e-01, -5.38096130e-01, -3.52915227e-01, -2.89237618e-01, + -4.53390926e-01, -3.57132286e-01, -2.32824370e-01, -1.73414394e-01, + -3.02039832e-01, -4.06891137e-01, -3.11633587e-01, -3.86790335e-01, + -4.44779336e-01, -4.03608739e-01, -3.27035785e-01, -5.77829838e-01, + -3.05345178e-01, -5.01044929e-01, -4.90038365e-01, -3.42383653e-01, + -2.16015473e-01, -2.34810472e-01, -3.94850314e-01, -3.71235818e-01, + 1.26175291e-03, -4.32831019e-01, -1.67813487e-02, -4.22147512e-01, + -3.73436451e-01, -4.40798849e-01, -4.10787642e-01, -4.72505748e-01, + -2.18411952e-01, -3.72730881e-01, -3.14533114e-01, -2.04394564e-01, + -2.78175622e-01, -2.42158204e-01, -3.91372234e-01, -5.00931203e-01, + -2.58081257e-01, -4.08814311e-01, -4.02594358e-02, -3.35627794e-02, + -2.13214949e-01, -3.36978287e-01, -2.62033820e-01, -6.90276861e-01, + -3.95373791e-01, -4.39152598e-01, -2.85017580e-01, -4.47889686e-01, + -4.36591983e-01, -3.10838968e-01, -4.41035658e-01, -2.79656291e-01, + -1.40883625e-01, -3.65119666e-01, -1.61859736e-01, -1.76689908e-01, + -3.66094023e-01, -4.24499989e-01, -4.19989884e-01, -2.08415791e-01, + -2.32176647e-01, -3.55920881e-01, -2.76795477e-01, -4.64629948e-01, + -2.86160856e-01, -5.59463799e-01, -3.76166940e-01, -3.82688463e-01, + -4.46783721e-01, -3.94609421e-01, -3.84682864e-01, -2.80228592e-02, + -3.13305557e-01, -3.56978267e-01, -4.11862016e-01, -6.72479212e-01, + -2.63180643e-01, -2.53385037e-01, -4.69251305e-01, -4.94098127e-01, + -3.67969900e-01, -5.01329899e-01, -4.63324428e-01, -4.17908072e-01, + -2.42769625e-03, -3.66800427e-01, -4.42783326e-01, -5.09920537e-01, + -3.79778743e-01, -2.11185992e-01, -4.76333320e-01, -2.97944933e-01, + -5.35430849e-01, -5.54302871e-01, -1.78685069e-01, -5.05887330e-01, + -3.60490650e-01, -2.97754109e-02, -4.16262656e-01, -2.02613860e-01, + -4.54910725e-01, -5.27239859e-01, -3.19837511e-01, -3.31910402e-01, + -5.04021525e-01, -5.96606016e-01, -2.31208146e-01, -5.22311389e-01, + -3.61870646e-01, -2.52372891e-01, -5.43789089e-01, -3.68237197e-01, + -2.00233698e-01, -3.31239879e-01, -3.89182925e-01, -1.77743852e-01], dtype=float32), + kernel: Array([[-0.03629679, 0.12771748, 0.30917847, ..., -0.02645334, + 0.2853967 , 0.25290728], + [-0.05366539, -0.08199085, -0.77184826, ..., -0.28385365, + -0.08720582, -0.2480316 ], + [-0.30405623, -0.14440256, -0.16176778, ..., -0.6061087 , + -0.53592414, -0.557683 ], + ..., + [-0.5002675 , 0.6923438 , -0.39955443, ..., 0.2793648 , + -0.0890847 , 0.13734408], + [ 0.42013338, -0.1837182 , 0.20519936, ..., 0.59244996, + 0.15519385, -0.04555327], + [ 0.2887559 , 0.03648978, 0.19605577, ..., -0.3638592 , + -0.05200996, 0.14020972]], dtype=float32), + }, + output: { + bias: Array([-2.03454956e-01, 2.89009154e-01, -5.26833301e-03, 1.99230332e-02, + 7.02989474e-03, -8.10585693e-02, 1.01039574e-01, 1.28230199e-01, + 1.65420137e-02, 1.76779047e-01, 2.31052637e-02, -2.84619331e-02, + -1.72500789e-01, 8.57431665e-02, -2.58707665e-02, 8.68742466e-02, + 3.48051847e-03, -8.32902640e-02, 8.47522356e-03, -6.26579765e-03, + 6.64773211e-02, -5.39215617e-02, -1.33164413e-02, -2.61053070e-02, + -1.25059590e-01, 9.69652552e-03, -6.74541369e-02, 5.93836829e-02, + 6.43064156e-02, 6.59576431e-02, 1.18919395e-01, -1.21684916e-01, + 1.34139672e-01, -7.79727027e-02, 7.02831373e-02, -7.30708465e-02, + 2.56273188e-02, 8.22603703e-03, 5.56182116e-03, -1.27773434e-01, + -1.63553078e-02, -2.40529538e-03, -7.88959324e-01, 2.98491977e-02, + -1.15960285e-01, -1.11549925e-02, 1.14158385e-01, 4.73596305e-02, + -4.17886041e-02, -5.26564419e-02, -9.78346448e-03, 4.26121093e-02, + -1.87094569e-01, 8.54555890e-02, -4.69494089e-02, -3.20936963e-02, + -1.27747595e-01, -3.98493744e-02, 1.56673312e-01, 5.93540259e-02, + 1.29704997e-01, -2.26318836e-01, -4.20836620e-02, 1.89136285e-02, + 2.25093830e-02, -1.10592782e-01, -8.14836994e-02, -1.74993854e-02, + 6.80994093e-02, -1.73572097e-02, -1.06341034e-01, 2.98510771e-03, + -6.83057960e-03, -4.16511074e-02, 8.99981439e-01, 1.04607493e-02, + -1.10422350e-01, 4.36907373e-02, -1.15947202e-01, -4.37050452e-03, + -4.16922748e-01, 3.20061632e-02, 4.93779741e-02, 2.31584487e-03, + -1.06097618e-02, 8.79805684e-02, 1.83927685e-01, 1.49958774e-01, + 1.10785790e-01, 1.96547806e-02, -6.48459792e-03, 2.37252992e-02, + -1.06167868e-01, -1.58290878e-01, -3.44699137e-02, -5.64344190e-02, + 2.31592474e-03, 1.24252474e-04, -6.33904487e-02, 1.05357006e-01, + 2.25357682e-01, -6.24567755e-02, 3.09208464e-02, -1.09267943e-01, + 1.43063843e-01, -6.33983240e-02, -1.62469558e-02, -1.03259414e-01, + -8.94914567e-02, -3.48297022e-02, -1.44557551e-01, 1.30005326e-04, + 1.44872010e-01, -1.74515136e-02, 4.52921912e-02, 1.08484700e-01, + -9.68603417e-03, -4.61588614e-03, -3.65241021e-02, -3.43756750e-04, + 1.57616019e-01, 8.44166130e-02, 2.49973107e-02, -9.32366401e-02, + 2.59006135e-02, 4.68501188e-02, 5.38958088e-02, 1.02374971e-01], dtype=float32), + kernel: Array([[-0.35964 , 0.58528113, 0.12339137, ..., -0.2728965 , + -0.35934824, -0.06503713], + [ 0.35246447, -0.02949085, 0.5403422 , ..., -0.09529549, + 0.11235058, -0.22131431], + [ 0.13894935, -0.3717892 , -0.66237545, ..., -0.41706035, + 0.52529556, 0.16840486], + ..., + [-0.09083151, 0.03931169, 0.2619683 , ..., -0.20559356, + 0.15253511, 0.3468645 ], + [-0.02631908, -0.0662327 , 0.12175545, ..., 0.10028415, + -0.3420218 , 0.01938576], + [-0.15720125, 0.16154028, -0.15429054, ..., -0.42136976, + 0.29129225, 0.06028228]], dtype=float32), + }, + }, + }, + }, + 2: { + attention: { + attention: { + key: { + bias: Array([ 4.91330713e-01, 9.67662875e-03, 5.60255200e-02, 8.97967637e-01, + -3.06834251e-01, 1.06595111e+00, 5.69995821e-01, -9.72940326e-01, + -1.00492425e-01, 2.51781523e-01, -3.73477489e-01, 1.62789464e-01, + -3.40229303e-01, 8.37876360e-05, 4.15999711e-01, 8.95448923e-02, + -1.16417743e-01, 1.02886930e-01, -1.22432733e+00, -2.46076301e-01, + -5.80214441e-01, -4.96384725e-02, -2.68343747e-01, -3.29451412e-01, + 2.24202454e-01, -6.01632036e-02, -6.19756997e-01, -3.03309470e-01, + -1.28122672e-01, -2.58336246e-01, 4.15848672e-01, 6.79572225e-02, + 1.68727383e-01, -1.08576640e-01, 9.50729072e-01, 1.18212245e-01, + 2.77229160e-01, 4.02590066e-01, 2.45219976e-01, -2.20667735e-01, + -6.55377328e-01, -2.22052753e-01, -2.34352529e-01, 7.21003652e-01, + -2.47502819e-01, -8.39031518e-01, 7.14816153e-01, 5.66149652e-01, + 1.96338549e-01, -3.21684837e-01, 1.26060262e-01, -1.13100424e-01, + 7.06214011e-02, -3.80029410e-01, -5.25692463e-01, 5.33015609e-01, + -2.28610292e-01, 8.77432287e-01, 4.54767972e-01, 1.05058357e-01, + -1.01594612e-01, -4.25737798e-01, 8.72221589e-01, -4.27799642e-01, + -7.26574242e-01, 3.97354126e-01, 5.40428281e-01, -6.92076981e-01, + -7.33269453e-01, -8.38262439e-01, -9.10226703e-01, -2.79593736e-01, + 5.23640335e-01, 2.84005404e-01, 3.39674622e-01, -2.09556028e-01, + -1.83250681e-02, 2.20996439e-01, 8.39932729e-03, -1.35142893e-01, + -4.47591990e-01, -5.44816703e-02, 3.56293529e-01, 3.84467632e-01, + -4.40692216e-01, 1.00856110e-01, 8.52475286e-01, 3.95925105e-01, + -5.06571472e-01, 4.37813222e-01, -8.72387230e-01, -4.80205119e-01, + -6.85112596e-01, -6.57971025e-01, 3.89101595e-01, -2.59548306e-01, + -2.97444463e-01, 2.12126691e-02, -8.14053118e-02, 1.06517822e-01, + 2.81973779e-01, 7.66788945e-02, -4.01952624e-01, -1.77236781e-01, + 6.38018996e-02, -3.04998636e-01, -3.08628559e-01, 1.56093314e-01, + 1.61672279e-01, 1.33270577e-01, 1.59121871e-01, -1.46263152e-01, + 1.54846618e-02, 2.81419247e-01, -3.13295014e-02, 2.87094891e-01, + -3.95765811e-01, 4.69570488e-01, -4.26991880e-01, 2.32855052e-01, + -1.62792206e-01, -3.81025262e-02, 4.74744802e-03, -7.95326605e-02, + 1.28198788e-02, 1.01442777e-01, -4.87101451e-02, -3.22961807e-02], dtype=float32), + kernel: Array([[-0.20145601, -0.22677487, -0.05235473, ..., -0.06058222, + -0.01446399, 0.16699584], + [ 0.22185628, 0.15144752, 0.23851915, ..., -0.10337625, + -0.07871008, 0.01390354], + [ 0.01399281, 0.06500424, -0.15749013, ..., 0.14727156, + -0.08356368, 0.03307934], + ..., + [ 0.01401845, -0.0230756 , -0.0378951 , ..., -0.40631384, + -0.10683019, 0.05925401], + [-0.00425362, -0.07440671, 0.03620811, ..., 0.21968801, + -0.18367934, 0.02989441], + [-0.1020968 , 0.5364823 , 0.00217419, ..., -0.05559423, + 0.13003536, -0.06085972]], dtype=float32), + }, + query: { + bias: Array([ 0.3528763 , 0.3471857 , 0.21516396, 0.44737092, -0.4285367 , + 0.3081475 , -0.25559098, -0.5126486 , 0.21324182, -0.32549337, + 0.02813848, -0.38898697, -0.34589097, -0.19787419, 0.28709203, + -0.24419545, -0.26928642, 0.30525193, -0.32151905, 0.0448105 , + 0.51535946, -0.11308531, 0.3212626 , 0.4922951 , 0.26803052, + -0.38731664, -0.2744102 , 0.39860168, 0.4643956 , 0.20418273, + 0.14548776, 0.0657926 , -0.21296258, 0.15750653, -0.3648503 , + 0.61559546, -0.35865372, 0.0193048 , 0.31175947, 0.2440965 , + 0.318281 , -0.7344339 , -0.88320374, -0.05781611, 0.18718946, + 0.20290868, -0.23830453, 0.5799299 , 0.01096278, -0.4863952 , + -0.03866521, 0.30802044, 0.44045374, -0.03443577, -0.02167755, + -0.01554048, -0.511046 , 0.41503018, -0.58751154, 0.5616692 , + -0.26840666, -0.20692529, 0.71333766, 0.2290065 , 0.15057702, + 0.1327177 , 0.31366843, 0.3856127 , -0.18231715, -0.24914798, + 0.321846 , -0.3574352 , -0.12077462, 0.36877826, 0.38871482, + -0.3660267 , 0.7790568 , 0.06110919, -0.22040306, -0.23536305, + -0.07361323, -0.47640693, 0.53327733, 0.49116883, -0.5596567 , + -0.7786464 , -0.1766858 , -0.14165151, -0.4430259 , 0.38748688, + 0.02364437, -0.46849003, -0.23776187, -0.45196384, 0.65864664, + 0.39194986, 0.07788196, 0.10206337, 0.15233228, 0.01135532, + 0.08919958, -0.0229385 , -0.00552564, -0.6363855 , 0.04719624, + 0.07559558, 0.0246234 , 0.32256266, -0.02410899, -0.37407702, + -0.063107 , -0.24087183, -0.0432428 , 0.01937803, -0.12409952, + 0.38090006, -0.8962724 , 0.51235485, -0.21019219, -0.06553419, + -0.0576391 , -0.00670896, 0.0239524 , 0.00629916, 0.46627957, + 0.02785698, 0.03592508, 0.06080098], dtype=float32), + kernel: Array([[ 0.09377078, 0.06110698, -0.01986536, ..., -0.09585967, + -0.00882398, -0.09038808], + [-0.1099542 , -0.0294884 , -0.3200316 , ..., 0.14632395, + -0.085163 , -0.12561767], + [-0.17357844, -0.08655021, -0.17681202, ..., 0.2461566 , + -0.11222455, 0.1451854 ], + ..., + [ 0.01483 , 0.12659054, 0.3020677 , ..., -0.1301052 , + -0.02523333, 0.0047435 ], + [-0.21112014, -0.10351861, 0.12283702, ..., 0.04758199, + -0.02335739, -0.20571615], + [ 0.04751823, -0.08418679, 0.021736 , ..., 0.23970562, + -0.10594929, -0.2381722 ]], dtype=float32), + }, + value: { + bias: Array([-0.20210104, -0.22905709, -0.14039448, -0.03683947, 0.12594987, + -0.02302479, 0.00906336, -0.00965133, -0.07539421, 0.19282112, + -0.14035326, -0.02629162, -0.0613457 , 0.19251512, 0.22437336, + 0.01865598, 0.09472503, 0.07302232, 0.23662573, 0.1634112 , + -0.25386894, 0.03644704, -0.14476314, -0.15108086, 0.30487335, + -0.01728077, -0.03989496, 0.10778196, 0.01612386, -0.02545331, + 0.10517047, 0.00122259, 0.06249418, 0.1109454 , 0.07476445, + -0.00254933, -0.08234914, 0.00984051, -0.18144692, 0.16224049, + 0.0791403 , 0.03824679, 0.05325351, 0.0459534 , -0.01707704, + -0.05914249, -0.07880631, 0.09950723, 0.01833864, -0.12028188, + 0.0201016 , 0.17571941, 0.08609623, -0.07171538, 0.27483281, + 0.19407238, 0.01660765, 0.03474303, 0.09089345, 0.19815543, + 0.21818227, 0.1436114 , 0.00318472, 0.06330492, -0.06290373, + 0.03819067, -0.06003238, -0.03846627, 0.07632727, -0.05140862, + -0.07411028, 0.21224949, -0.07434353, -0.09569198, -0.01517388, + 0.00378666, 0.0954181 , 0.02467517, -0.0480962 , -0.07227837, + 0.08442853, -0.11686811, 0.04613666, -0.04021059, 0.03575824, + 0.02964796, 0.09947117, 0.07083944, -0.06642204, -0.07676494, + 0.02942543, -0.0550985 , -0.06826791, 0.06714635, -0.05162644, + -0.03433613, -0.00067884, 0.05411651, -0.01232108, 0.10868998, + 0.05962406, -0.00136539, 0.01625578, -0.00632739, -0.0961687 , + -0.07040943, 0.02730714, 0.00544707, 0.01023827, 0.03284573, + 0.02555647, 0.0085297 , 0.03900068, 0.02747485, -0.03433363, + -0.01524182, -0.00976253, -0.06111204, 0.10030263, 0.04285727, + -0.00326952, -0.04434157, 0.11089424, 0.04277305, 0.02115745, + 0.02765508, 0.08399435, 0.08146379], dtype=float32), + kernel: Array([[-1.31283207e-02, 5.85139990e-02, 6.22775187e-05, ..., + 4.85775769e-02, -4.02368233e-02, -6.80026598e-03], + [-2.73164004e-01, -3.41370434e-01, -2.39238873e-01, ..., + 4.36691791e-01, 3.35174829e-01, 8.86783656e-03], + [ 6.34363815e-02, 3.22204083e-01, 2.63729244e-01, ..., + -9.32745263e-03, -1.88997626e-01, 2.92390492e-02], + ..., + [-1.07346885e-01, -1.71600416e-01, -1.54479533e-01, ..., + -2.66532987e-01, -7.03294575e-02, -3.63519728e-01], + [-1.77037895e-01, -1.31959140e-01, 3.59883085e-02, ..., + 1.12166613e-01, -3.25531811e-01, 4.15854841e-01], + [ 1.20714135e-01, 1.70483068e-01, 6.37787461e-01, ..., + -3.24400365e-01, -7.82848775e-01, -1.08967312e-01]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-1.42582014e-01, 3.51325125e-01, -5.49227633e-02, -1.36979753e-02, + -8.89001694e-03, -4.82381359e-02, 7.59436563e-02, 4.62903231e-02, + 1.25182196e-02, 1.74479246e-01, 5.47183445e-03, -4.21341211e-02, + -1.15092970e-01, 5.72893657e-02, -5.69992810e-02, 5.97054660e-02, + 2.31458675e-02, -9.98827890e-02, -1.63127463e-02, -5.68029238e-03, + 1.06309690e-01, -5.58492132e-02, -2.88810022e-02, -5.41146100e-02, + -1.25681654e-01, 1.55297220e-02, -1.05181850e-01, 2.84042470e-02, + 8.00593048e-02, 1.12994172e-01, 6.64846674e-02, -8.10639188e-02, + 6.96494877e-02, -7.10230693e-02, 2.48000082e-02, -6.40962422e-02, + 5.03854305e-02, 2.99889389e-02, 2.71755550e-02, -1.32961541e-01, + -9.42786504e-03, 6.68036118e-02, -7.92975008e-01, 2.92515382e-02, + -1.21656477e-01, 6.97143525e-02, 9.77309272e-02, 9.14023891e-02, + -1.06936611e-01, -4.60654832e-02, -3.59379910e-02, 6.73950538e-02, + -3.20891030e-02, 1.82818487e-01, -4.29846235e-02, 1.71033782e-03, + -2.07776949e-01, -8.03176984e-02, 1.26594633e-01, 1.53843045e-01, + 9.50956270e-02, -1.38820022e-01, -5.31725921e-02, 3.35701555e-02, + 5.83558269e-02, -1.29202589e-01, -4.29277681e-02, 5.83481155e-02, + 1.43310845e-01, 3.71531360e-02, -1.16031744e-01, 1.25720399e-03, + -3.66386138e-02, -5.68410717e-02, 8.81127954e-01, -3.54486227e-04, + -1.14806622e-01, 5.00139184e-02, -1.25719443e-01, -2.92788483e-02, + -5.60584605e-01, 1.90249607e-02, 7.62136094e-03, 3.29709910e-02, + -1.73091292e-02, 1.49499878e-01, 1.71786591e-01, 1.50196671e-01, + -4.66331327e-03, 7.90596679e-02, -6.07569050e-03, 1.27625940e-02, + -6.54113144e-02, -1.46687955e-01, -2.00468884e-03, -7.76169971e-02, + -5.61596686e-03, -2.51500830e-02, -5.23540825e-02, 1.42120197e-01, + 1.58276126e-01, -3.11181620e-02, 4.16391157e-02, -9.11307335e-02, + 1.01979673e-01, -1.40924305e-01, -2.65695825e-02, 3.79197043e-03, + -9.74692628e-02, -5.42117134e-02, -1.30008578e-01, 3.88463400e-02, + 1.63258657e-01, -1.50325231e-03, 3.45914699e-02, 1.69843420e-01, + 1.55994203e-02, -1.33964028e-02, -1.95804201e-02, 1.02751255e-02, + 1.84371233e-01, 7.10138083e-02, 5.72973862e-03, -1.13018960e-01, + -8.22554808e-03, 5.45389391e-02, 1.77288230e-03, 1.88416634e-02], dtype=float32), + kernel: Array([[ 0.15479796, 0.53571814, 0.51493394, ..., -0.80009085, + -0.05134518, 0.13081257], + [-0.00908494, 0.42273548, 0.710739 , ..., -0.3385195 , + -0.41763583, 0.50498915], + [ 0.21792167, 0.02829238, -0.15942883, ..., 0.14649649, + -0.3748477 , 0.5899068 ], + ..., + [ 0.69713473, -0.10411429, 0.20231156, ..., 0.13196032, + -0.09708058, -0.21109731], + [ 0.5224463 , -0.00511254, -0.9390049 , ..., 0.06903978, + 0.24796766, 0.5909974 ], + [-0.6673977 , 0.24624978, -0.7719052 , ..., -0.09019069, + -0.29529244, -0.12763341]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([-2.69522518e-02, 7.07723275e-02, -1.92383816e-03, -1.69017464e-01, + 6.28127530e-03, 1.35174885e-01, 1.43263070e-02, -1.49136512e-02, + 1.43651441e-02, 2.47675721e-02, 2.31170952e-01, -1.36354938e-02, + -5.00643514e-02, 4.64095622e-02, 7.30278119e-02, 2.58173402e-02, + -7.96374157e-02, -9.20170024e-02, -1.74471997e-02, 1.98375121e-01, + 2.97931284e-01, 1.71703979e-01, -1.74371853e-01, -8.47198628e-03, + -2.85758138e-01, 2.43080053e-02, -2.33928546e-01, 1.65304869e-01, + 1.08051397e-01, -3.64561914e-03, -4.06746566e-01, -2.43781097e-02, + -1.94358621e-02, 2.50640996e-02, 1.21675162e-02, 3.38391075e-03, + 4.25222097e-03, -1.16181098e-01, -6.87912405e-02, -1.43533750e-02, + 6.30778596e-02, 5.16471975e-02, 2.78629586e-02, -6.81551844e-02, + -2.16658652e-01, 1.19643852e-01, -1.77933127e-02, 4.19648476e-02, + -6.94157332e-02, 8.80410895e-02, 3.07484251e-03, -2.26102751e-02, + 2.71407124e-02, -1.37352645e-02, 1.24339305e-03, -1.10182655e-03, + -4.25179452e-02, 3.90366204e-02, -6.09224401e-02, 7.32898619e-03, + -5.39776146e-01, 8.10408741e-02, 1.43810920e-02, -1.46246385e-02, + 1.65752345e-03, -2.01134533e-01, 1.59140024e-02, 4.93869092e-03, + 4.72723007e-01, -1.37103930e-01, 1.31060192e-02, -4.23336178e-02, + 4.09131646e-02, 3.04093435e-02, 5.28761335e-02, 3.03985029e-01, + -2.57239580e-01, 1.94385331e-02, 2.98021808e-02, -3.17842066e-02, + -3.57343435e-01, 6.69872388e-02, -2.05376260e-02, 6.03192970e-02, + -2.77649582e-04, 4.66996022e-02, 2.09306199e-02, 1.26473635e-01, + -1.99992478e-01, 6.02058368e-03, -9.60124508e-02, -6.86570778e-02, + -2.69254763e-02, -1.16648778e-01, 6.80445060e-02, 8.54630917e-02, + 6.16764985e-02, -6.65853033e-03, -6.19218759e-02, -3.73522252e-01, + -1.32694215e-01, 4.03534174e-02, 5.64134538e-01, 5.32403849e-02, + -7.38940388e-02, -9.58362669e-02, -1.27810821e-01, 5.41193448e-02, + 1.32223934e-01, 4.17461758e-03, 7.13659152e-02, 3.29776555e-02, + 3.19321036e-01, -3.77637222e-02, 2.86701359e-02, 9.89994928e-02, + 5.24967327e-04, 6.47896230e-02, -9.30845514e-02, -8.97478983e-02, + 3.75827640e-01, 7.06133619e-03, -1.70939237e-01, -2.22015311e-03, + -8.06940794e-02, -1.80535465e-01, -1.17046207e-01, -5.05085103e-02], dtype=float32), + scale: Array([-2.16153208e-02, 4.99568492e-01, 1.90679863e-01, 3.98646802e-01, + 3.23043484e-03, 8.04649472e-01, 1.73618346e-01, 6.32202104e-02, + 6.03341032e-03, 7.60933161e-01, 5.64841628e-01, 3.30200821e-01, + -1.13521464e-01, 4.35642749e-01, 7.02156186e-01, 1.86426133e-01, + 2.69546956e-01, 4.07335073e-01, 3.06657463e-01, 7.59951055e-01, + 8.35915983e-01, 5.73645175e-01, 3.52393031e-01, 2.92183794e-02, + 5.95538437e-01, -3.08835153e-02, 5.68952084e-01, 6.18548512e-01, + 7.05874801e-01, -1.11141577e-01, 7.34429657e-01, 1.59030482e-01, + 6.69286311e-01, 3.07584852e-01, 2.97907770e-01, 2.45770905e-02, + 3.42297733e-01, 6.38300121e-01, 4.57696557e-01, 3.39488059e-01, + 4.27823365e-01, 6.61583364e-01, 2.51028594e-02, 6.87740266e-01, + 6.38814211e-01, 6.12671137e-01, 3.40550840e-01, 3.15654457e-01, + 5.78372836e-01, 5.64573705e-01, 5.55405932e-05, 3.71645719e-01, + 8.48703533e-02, 9.91551802e-02, 4.97666523e-02, 5.78866065e-01, + -1.33863851e-01, 2.80352771e-01, 5.00911117e-01, 1.85352847e-01, + 7.25611448e-01, 4.00964379e-01, 2.84439921e-01, 3.20623189e-01, + -1.70927703e-01, 4.05740440e-01, 3.79946142e-01, 2.70140201e-01, + 5.38314104e-01, 6.14893734e-01, 5.95924675e-01, 5.98044157e-01, + 1.48603588e-01, 3.81722450e-01, 5.22187948e-01, 7.07299054e-01, + 4.85696197e-01, 1.27313018e-01, 7.40394592e-02, 3.15272629e-01, + 3.43401968e-01, 6.31168902e-01, 2.14417100e-01, -3.38677377e-01, + 2.39464089e-01, 4.71792251e-01, 3.14170390e-01, 8.54388118e-01, + 6.38605714e-01, 1.02670453e-01, 4.03134346e-01, 3.42762381e-01, + 4.40316945e-01, 5.73590755e-01, 3.09652478e-01, 1.55366600e-01, + 6.81986809e-01, -7.89772999e-03, 5.20319641e-01, 6.24004126e-01, + 5.56326687e-01, 3.05760831e-01, 9.59146619e-01, 5.92343211e-01, + 3.85424703e-01, 4.82657462e-01, 5.25808513e-01, 4.87956852e-01, + 3.75198275e-01, -6.13166541e-02, 5.11129975e-01, 6.06624246e-01, + 7.52410650e-01, 1.93235308e-01, -5.42348363e-02, 3.70484799e-01, + 6.81148842e-03, 7.82989919e-01, 6.00687087e-01, 4.27850395e-01, + 9.20885563e-01, 2.06756860e-01, 7.61731803e-01, 6.10257149e-01, + 2.84399986e-01, 6.83825493e-01, 5.32033622e-01, 1.40386716e-01], dtype=float32), + }, + layernorm_before: { + bias: Array([ 1.7358573e-02, -2.1321216e-01, -3.0296935e-02, -8.5394399e-04, + -3.9229371e-02, 5.5817175e-03, -2.7252757e-03, -4.4029057e-02, + -2.0724675e-03, -1.2441806e-01, -2.2284750e-02, -2.7261129e-02, + -2.7341412e-03, 3.6942791e-03, 4.5538493e-03, 7.4563054e-03, + -6.2132953e-03, 9.7167321e-02, -2.5345799e-02, -3.9951671e-02, + -4.3273875e-03, -1.4203915e-03, 6.3152872e-02, 1.3384975e-03, + -5.5942591e-02, 6.3980678e-03, -3.2894940e-03, 2.6343842e-03, + 2.8435083e-02, 2.8260322e-02, 1.2178370e-02, -5.3195474e-03, + -2.4072591e-03, 4.1922098e-03, -7.3268535e-03, 7.5856894e-03, + 3.8502638e-03, -9.8595031e-02, 2.4951894e-02, -9.5122019e-03, + -5.4523430e-04, 4.9118581e-03, 2.0277210e-01, -1.6341466e-02, + 3.9525043e-02, -2.2266436e-02, -2.0717401e-03, -1.2440056e-01, + 2.4503253e-03, 1.1633219e-02, 3.5871702e-04, 1.0992530e-03, + 7.6659650e-02, -6.7069583e-02, 9.7138679e-04, -2.9395521e-02, + -8.4895557e-03, -1.3592694e-02, 2.9208256e-02, 1.7238649e-02, + 5.7587404e-02, 4.4692371e-02, 1.6647557e-02, -6.6133849e-03, + -3.3748765e-03, 9.0554915e-03, 1.8321996e-03, 1.5999395e-02, + -1.3479884e-02, 2.9679830e-03, -2.7784597e-02, 3.4300985e-03, + 4.0611736e-02, -4.6936791e-02, -7.2169447e-01, 8.0935545e-03, + -1.8306768e-02, -5.0263423e-02, -3.8130886e-03, 1.9901623e-03, + 2.5550249e-01, 6.1443527e-03, -2.2893531e-02, -1.7245837e-02, + 2.9439591e-02, -4.3969110e-02, -6.7670122e-03, -1.3685081e-02, + -9.9776837e-04, 9.0234309e-02, -1.4077160e-02, 2.6768299e-03, + -3.8174929e-03, -1.4334331e-02, 3.4874912e-02, 8.0644749e-02, + 8.6499676e-03, -1.7143000e-02, 7.9911202e-03, 2.8301086e-02, + -4.1625272e-02, -3.0587490e-03, -5.6571878e-02, 9.0596341e-03, + -7.6397270e-02, -4.5915369e-02, -6.2590851e-03, 3.4895688e-02, + 6.1382476e-02, -5.6433260e-02, 1.0131227e-02, 1.5290984e-03, + -5.1695153e-02, 1.3336526e-01, 1.1029670e-02, -2.2941723e-02, + 2.4411927e-03, -2.4650490e-02, -1.1010250e-02, -7.5843465e-03, + -1.1662335e-02, 3.2463686e-03, -1.6854508e-03, 2.7790260e-02, + 7.3739677e-03, -9.9853233e-02, 1.0033315e-02, -9.9916525e-02], dtype=float32), + scale: Array([ 7.94857554e-03, 1.53508723e-01, 7.40691870e-02, 1.35285795e-01, + 2.26536706e-01, 1.69742465e-01, 8.65431502e-04, 3.19571525e-01, + 1.62680820e-03, 1.62808627e-01, 1.34147704e-01, 2.01822311e-01, + 6.27691345e-03, 7.19762873e-03, 3.74700986e-02, 3.44181597e-01, + 9.95386466e-02, 1.60152465e-01, 9.15851742e-02, 1.63604781e-01, + 1.17830604e-01, -1.66135898e-03, 2.23148689e-01, 1.00941835e-02, + 2.60033756e-01, 1.75621884e-03, 2.65612965e-03, 4.38557938e-04, + 3.29755515e-01, 2.43666410e-01, 3.20717841e-01, 1.77067131e-01, + 2.68771589e-01, 3.02107055e-02, 2.95764953e-02, 1.08823786e-02, + -1.60506088e-02, 2.95251936e-01, -1.78114280e-01, 8.72005969e-02, + -8.84137396e-03, 2.99656868e-01, 8.63162801e-02, 1.33525252e-01, + 3.03551525e-01, 2.60979742e-01, -5.50631993e-03, 2.83051997e-01, + 2.56603450e-01, 2.00551033e-01, 5.14968205e-03, 2.78209685e-04, + 2.23627791e-01, 2.34901443e-01, -8.34921002e-03, 2.24900410e-01, + -6.23680511e-03, -2.86224540e-02, 2.08953783e-01, 3.78563493e-01, + 2.40491629e-01, 1.92855686e-01, 2.62306392e-01, 5.38155213e-02, + 2.10460991e-01, 1.67684779e-02, 5.48215546e-02, 1.18601710e-01, + 3.43158811e-01, 2.41343841e-01, 2.74241388e-01, -6.84597194e-02, + 2.30040535e-01, 2.19160318e-01, 1.71067789e-01, 2.74585098e-01, + 1.48560002e-01, 2.04053581e-01, -1.04316017e-02, 3.06529433e-01, + 1.51589945e-01, -2.79646995e-03, 2.60486454e-01, 9.33720767e-02, + 2.20337927e-01, 1.95606187e-01, -4.61113850e-05, 7.93705583e-02, + 2.31670171e-01, 1.81787178e-01, 6.02688640e-02, 6.66163536e-03, + 5.13362838e-03, 8.27364624e-02, 3.05262417e-01, 1.02641329e-01, + 1.57344952e-01, 1.41675845e-01, 4.87802178e-03, 1.00217178e-01, + 1.87011242e-01, -1.03036012e-03, 2.17317730e-01, 1.16062611e-01, + 3.36587638e-01, 3.04190874e-01, 1.32016614e-01, 1.66113153e-01, + 3.42909336e-01, 2.46614054e-01, 2.17617258e-01, 1.86562017e-01, + 1.04256548e-01, 3.53437752e-01, 5.86070295e-04, 2.16139719e-01, + 2.78897047e-01, 2.12332204e-01, 3.80921900e-01, 1.19425990e-01, + 3.23329829e-02, -1.03264116e-02, 1.89511199e-02, 1.55454144e-01, + -3.36376615e-02, 2.82643199e-01, 3.45727682e-01, 2.41249815e-01], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-2.87919283e-01, -1.59886956e-01, -3.92205477e-01, -4.32191879e-01, + -3.32138002e-01, -5.40607274e-01, -3.54391366e-01, -1.31854281e-01, + -5.39756715e-01, -4.69244629e-01, -3.44658494e-01, -4.88029271e-01, + -3.98942500e-01, -4.15098637e-01, -3.66716921e-01, -5.07466793e-01, + -2.69494671e-02, -3.29423159e-01, -3.49476814e-01, -1.71700358e-01, + -3.80648106e-01, -4.23726380e-01, -4.98886943e-01, -4.94191587e-01, + -2.60305732e-01, -3.48132432e-01, -4.27568406e-01, -3.35825920e-01, + -3.69937718e-01, -3.21786374e-01, -2.29871064e-01, -2.88744688e-01, + -3.85481507e-01, -4.31343943e-01, -1.41554534e-01, -1.79118499e-01, + -4.28618819e-01, -2.48103037e-01, -9.70930979e-02, -4.53777105e-01, + -1.97259694e-01, -3.89665425e-01, -4.33900356e-01, -2.01605186e-01, + -3.11053962e-01, -3.33365947e-01, -4.72573489e-01, -4.81180519e-01, + -4.59191829e-01, -3.63978118e-01, -5.11594951e-01, -3.21153313e-01, + -5.05568504e-01, -4.40596819e-01, -2.05569610e-01, -2.04613134e-01, + -2.30359152e-01, -4.58120763e-01, -1.35851100e-01, -2.78527498e-01, + -5.32928765e-01, -4.59086925e-01, -1.91062093e-01, -2.50361383e-01, + -3.23409230e-01, -5.09010673e-01, -2.47201502e-01, -3.10001701e-01, + -3.00568342e-01, -3.73322159e-01, -5.35207331e-01, -5.31656861e-01, + -3.18275809e-01, -4.40275580e-01, -3.43852162e-01, -4.12837327e-01, + -3.85347039e-01, -4.52745438e-01, -2.02980176e-01, -4.54181701e-01, + -5.00234187e-01, -3.34358454e-01, -3.39169592e-01, -2.89554566e-01, + -3.29827011e-01, -5.80846727e-01, -3.28373313e-01, -4.26805854e-01, + -3.15485120e-01, -3.92600238e-01, -3.12504858e-01, -2.19271436e-01, + -1.64767921e-01, -5.02310932e-01, -2.58719236e-01, -4.32749659e-01, + -9.22016874e-02, -4.70086277e-01, -3.97373617e-01, -4.89154935e-01, + -4.44214731e-01, -2.03317106e-01, -4.75122511e-01, -2.69819319e-01, + -2.67087281e-01, -4.90763247e-01, -4.22214806e-01, -2.64238060e-01, + -2.23214269e-01, -3.43690515e-01, -2.78915852e-01, -4.24799353e-01, + -5.69192111e-01, -2.07584277e-01, -1.64182380e-01, -1.91950992e-01, + -5.61932206e-01, -5.16008019e-01, -3.33474815e-01, -4.25828815e-01, + -1.61096185e-01, -3.98274213e-01, -5.68151116e-01, -2.51645744e-01, + -2.21852303e-01, -4.22558784e-01, -3.82831007e-01, -3.19712460e-01, + -1.66578472e-01, -2.49394879e-01, -4.19118315e-01, -2.93457031e-01, + -2.79008389e-01, -5.07268071e-01, -1.48156717e-01, -4.60653484e-01, + -2.05382407e-01, -4.70562130e-01, -2.83998400e-01, -4.29632753e-01, + -2.36523062e-01, -3.75578731e-01, -3.47366124e-01, -2.43162408e-01, + -4.38231677e-01, -5.36709070e-01, -1.51090562e-01, -4.71059680e-01, + -4.89868373e-01, -3.91524017e-01, -3.24053317e-01, -3.04224849e-01, + -4.36558515e-01, -2.35701188e-01, -4.07339513e-01, -2.95855463e-01, + -2.55902320e-01, -4.83437032e-01, -3.88129205e-01, -1.03971243e-01, + -2.60267794e-01, -3.92015398e-01, -4.49646682e-01, -2.48017311e-01, + -2.56850272e-01, -1.90488473e-01, -4.68013018e-01, -4.13734376e-01, + -5.12630165e-01, -4.26820993e-01, -3.74209106e-01, -3.59761298e-01, + -4.02847528e-01, -2.97854811e-01, -3.58279765e-01, -3.24418694e-01, + -3.66349548e-01, -4.33402061e-01, -4.91001993e-01, -5.03202379e-01, + -5.00787556e-01, -3.92910451e-01, -2.84466565e-01, -3.34319711e-01, + -4.25348043e-01, -2.26335734e-01, -2.37801537e-01, -4.87084776e-01, + -3.79190952e-01, -3.08594137e-01, -3.53578776e-01, -2.29536459e-01, + -4.85832751e-01, -3.31020027e-01, -3.49860758e-01, -4.98145282e-01, + -2.17634872e-01, -3.08456481e-01, -4.34388995e-01, -2.14153975e-01, + -2.20951349e-01, -4.75258082e-01, -2.87501752e-01, -9.52811018e-02, + -5.12147367e-01, -2.70634055e-01, -2.17771903e-01, -3.38476300e-01, + -3.06095690e-01, -3.33468020e-01, -1.24255188e-01, -5.16220808e-01, + -2.23946184e-01, -5.59429586e-01, -2.83430904e-01, -4.03296888e-01, + -3.68934721e-01, -3.25362712e-01, -7.39004537e-02, -3.09295535e-01, + -4.49250668e-01, -5.65693378e-01, -2.64678180e-01, -3.42923045e-01, + -3.37939441e-01, -3.71786445e-01, -4.06780869e-01, -5.13805926e-01, + -1.10545248e-01, -4.07462388e-01, -4.62072462e-01, -3.73685181e-01, + -4.82776046e-01, -1.83113769e-01, -1.37951612e-01, -4.32110876e-01, + -4.39051688e-01, -4.20039445e-01, -4.49939400e-01, -5.46136618e-01, + -4.43518162e-01, -3.36255163e-01, -2.60323614e-01, -1.48123205e-01, + -3.08105409e-01, -5.92286289e-01, -3.61411870e-01, -5.49996793e-01, + -4.72203493e-01, -1.12566963e-01, -3.46157908e-01, -2.21687108e-01, + -1.49518579e-01, -4.22669083e-01, -5.09057462e-01, -3.78608346e-01, + -2.69701719e-01, -2.57610738e-01, -3.43267351e-01, -3.99610460e-01, + -5.19208014e-01, -3.55502129e-01, -3.64218682e-01, -4.58064228e-01, + -1.66265354e-01, -4.56624329e-01, -4.21405643e-01, -3.38668913e-01, + -4.08919096e-01, -2.49855325e-01, -2.68893301e-01, 5.37731161e-04, + -4.11918223e-01, -2.53745079e-01, -2.37736523e-01, -3.77250522e-01, + -3.40394109e-01, -4.99163866e-01, -3.26041222e-01, -3.24631035e-02, + -4.25087065e-01, -3.99343103e-01, -4.45631206e-01, -2.19975799e-01, + -4.91138816e-01, -3.16317588e-01, -1.97857752e-01, -4.38183516e-01, + -2.38349304e-01, -3.12527388e-01, -4.46206242e-01, -1.97425947e-01, + -3.54068547e-01, -5.62409818e-01, -4.79922563e-01, -4.07248676e-01, + -5.08961380e-01, -3.35608751e-01, -4.69772518e-01, -4.64478105e-01, + -2.38058373e-01, -3.87834966e-01, -4.91571218e-01, -1.16468966e-01, + -1.88293502e-01, -2.19764918e-01, -2.41493329e-01, -3.08782160e-01, + -3.80097032e-01, -2.52258271e-01, -3.71472836e-01, -3.24205458e-01, + -4.05035138e-01, -5.40303230e-01, -1.73389211e-01, -3.20962131e-01, + -2.16032237e-01, -1.91938236e-01, -1.65382430e-01, -4.07554179e-01, + -3.40726435e-01, -2.95837790e-01, -9.07381624e-02, -1.67001665e-01, + -3.45026761e-01, -4.43528742e-01, -4.61021096e-01, -2.25696385e-01, + -2.09777251e-01, -5.53110957e-01, -3.23776424e-01, -3.86852056e-01, + -5.13776362e-01, -1.70138955e-01, -4.78795141e-01, -2.65858293e-01, + -1.22696631e-01, -5.30045271e-01, -3.90655428e-01, -2.64281482e-01, + -4.28456664e-01, -1.40250340e-01, -2.27036089e-01, -3.66672546e-01, + -4.42133546e-01, -4.25859839e-01, -2.67484695e-01, -4.62734908e-01, + -3.37051451e-01, -4.50581521e-01, -7.01922104e-02, -3.47592622e-01, + -3.91412824e-01, -4.24011290e-01, -3.32346976e-01, -3.72510850e-01, + -3.09896857e-01, -1.88529536e-01, -3.71884912e-01, -2.54814714e-01, + -3.63058448e-01, -4.75363225e-01, -3.06162030e-01, -4.65719730e-01, + -1.81163147e-01, -4.85125124e-01, -3.95848036e-01, -4.33863193e-01, + -3.04795027e-01, -4.90597248e-01, -2.44484976e-01, -3.44704390e-01, + -3.45088750e-01, -3.29458416e-01, -2.94500917e-01, -2.62573928e-01, + -1.69639900e-01, -4.71560776e-01, -2.42599428e-01, -3.69580209e-01, + -3.42656642e-01, -3.11950177e-01, -1.34305000e-01, -1.57368377e-01, + -3.91964465e-01, -3.12808394e-01, -2.37698436e-01, -3.69058728e-01, + -1.38768077e-01, -2.69422363e-02, -4.55560267e-01, -3.79797220e-01, + -4.36610907e-01, -2.92106032e-01, -3.98713797e-01, -4.69899833e-01, + -1.97569534e-01, -3.96044046e-01, -1.28238350e-01, -3.10444951e-01, + -4.51549977e-01, -2.95563221e-01, -4.12408620e-01, -4.34255183e-01, + -2.34395146e-01, -4.32429850e-01, -2.51709074e-01, -3.16113442e-01, + -2.24395722e-01, -4.12795454e-01, -3.96394074e-01, -3.76580715e-01, + -2.59395570e-01, -5.12954056e-01, -3.32104653e-01, -4.36192185e-01, + -5.13575554e-01, -1.39632359e-01, -4.51251060e-01, -4.95827556e-01, + -2.47558579e-01, -3.28758836e-01, -4.00682449e-01, -1.53779313e-01, + -2.39558816e-01, -3.46421301e-01, -2.07403153e-01, -4.59828466e-01, + -4.13055509e-01, -5.19209445e-01, -3.49062324e-01, -3.61741424e-01, + -5.46588659e-01, -3.25105667e-01, -2.64346927e-01, -3.62703383e-01, + -3.54973465e-01, -3.57288867e-01, -3.40226203e-01, -1.87480614e-01, + -3.84213597e-01, -3.67905647e-01, -2.21579894e-01, -4.32286561e-01, + -3.49667877e-01, -2.67124563e-01, -2.31143355e-01, -2.09680378e-01, + -4.44781452e-01, -1.10374503e-01, -1.90934509e-01, -4.75128025e-01, + -4.90421832e-01, -3.84051919e-01, -3.26278806e-01, -3.82725775e-01, + -2.96725929e-01, -4.35144246e-01, -3.04508001e-01, -4.23557281e-01, + -3.32650095e-01, -2.34854251e-01, -2.51995057e-01, -3.14376324e-01, + -4.71209139e-01, -5.88116467e-01, -3.59949678e-01, -3.27794611e-01, + -5.37792504e-01, -2.25097418e-01, -5.30355632e-01, -5.33735216e-01, + -2.61757433e-01, -2.64203936e-01, -3.20991725e-01, -3.19024026e-01, + -4.41864610e-01, -1.51468694e-01, -6.96782321e-02, -4.49556112e-01, + -3.41252834e-01, -6.03692472e-01, -4.20485705e-01, -4.03241307e-01, + -4.20258135e-01, -4.89987731e-01, -2.22824812e-01, -1.52097628e-01, + -3.01070571e-01, -3.31862330e-01, -2.91463912e-01, -2.47648209e-01, + -4.18553919e-01, 8.17964971e-02, -4.42326665e-01, -4.14036870e-01, + -4.23043162e-01, -4.24916655e-01, -3.22944313e-01, -3.46353531e-01, + -3.03149909e-01, -4.78416532e-01, -3.12508672e-01, -3.28645647e-01, + -3.10102880e-01, 6.61208481e-02, -3.10083032e-01, -4.34432387e-01, + -2.68977493e-01, -3.87750000e-01, -3.06754440e-01, -2.68470585e-01], dtype=float32), + kernel: Array([[-8.63173604e-02, 5.54475226e-02, -7.44883567e-02, ..., + -1.60854563e-01, -1.92300782e-01, 1.49715483e-01], + [-2.55667627e-01, -5.76032937e-01, -2.72577375e-01, ..., + -1.66639566e-01, -9.29499939e-02, -8.04038703e-01], + [ 1.05862938e-01, -7.12434500e-02, -3.51943791e-01, ..., + -4.56090540e-01, 6.42790049e-02, 2.83889264e-01], + ..., + [ 3.32350917e-02, -2.16388881e-01, -4.15698439e-01, ..., + -2.22950518e-01, 1.03157982e-01, 4.05906420e-03], + [ 1.34180099e-01, 2.93415725e-01, -3.85918528e-01, ..., + 5.76932847e-01, -1.18126400e-01, 1.22173645e-01], + [-3.73824616e-04, 1.94809213e-02, 2.37895355e-01, ..., + -2.55678475e-01, -4.13769521e-02, 8.79469365e-02]], dtype=float32), + }, + output: { + bias: Array([-0.14310208, 0.30402932, -0.04476445, 0.05614336, -0.00636329, + -0.10514647, 0.06905492, 0.06197605, 0.02608861, 0.15183868, + -0.0424929 , -0.01462859, -0.1350344 , 0.02212298, -0.05124652, + 0.06123491, 0.03247906, -0.09240599, -0.01013029, -0.03223318, + -0.0887438 , -0.06346248, 0.02817548, -0.05015382, -0.05139951, + 0.01631817, -0.04831094, -0.02007524, 0.04097073, 0.1089523 , + 0.14585266, -0.05194474, 0.04915197, -0.09083902, 0.01939188, + -0.06052509, 0.05472679, 0.04450262, 0.03008708, -0.1406127 , + -0.0286117 , 0.04946226, -0.8201933 , 0.05301183, -0.06421784, + 0.03291068, 0.08303672, 0.10530804, -0.05531342, -0.09611814, + -0.02686962, 0.10026816, -0.01932072, 0.17867358, -0.03787487, + 0.01405966, -0.20247376, -0.07591344, 0.15851738, 0.16269532, + 0.17163317, -0.16734251, -0.05690756, 0.02852942, 0.06089685, + -0.04197973, -0.06455492, 0.07540473, -0.06004319, 0.08624789, + -0.1040978 , 0.00247931, -0.02840352, -0.04520692, 0.8720378 , + -0.06833719, -0.07094937, 0.05479927, -0.12080298, -0.0121475 , + -0.5797822 , -0.02564085, 0.03635592, 0.03610236, -0.00385189, + 0.1565094 , 0.14343542, 0.05017952, 0.00714095, 0.09029737, + 0.0294838 , 0.01819415, -0.04793838, -0.1418107 , -0.00120118, + -0.06547695, -0.03586074, -0.01915462, -0.03162472, 0.25057518, + 0.21378383, -0.03219284, -0.03374225, -0.07630365, 0.13850403, + -0.09364668, 0.02092841, -0.02086668, -0.11372983, -0.06087056, + -0.14589557, 0.03926003, 0.0446431 , 0.00222743, 0.03577771, + 0.14117669, 0.01693195, -0.06432896, 0.02443499, 0.04788267, + -0.13781525, 0.06813591, 0.10359713, -0.10075218, 0.02409258, + 0.05323493, 0.00397177, 0.03774352], dtype=float32), + kernel: Array([[ 0.16250458, -0.15818997, 0.3543968 , ..., -0.06277009, + -0.458476 , 0.21236774], + [ 0.17157121, -0.33716652, 0.17193943, ..., -0.07994235, + 0.03173731, -0.09947567], + [ 0.20364675, 0.04469377, -0.3887421 , ..., -0.24748199, + 0.29170632, -0.02271716], + ..., + [ 0.29266545, -0.02244417, 0.41813877, ..., -0.37108174, + -0.26156428, -0.10771853], + [-0.07057951, -0.23378706, 0.05101302, ..., 0.08562624, + 0.45771068, 0.08117881], + [-0.21127233, -0.6031122 , -0.4248507 , ..., 0.16734137, + 0.09727449, -0.12857328]], dtype=float32), + }, + }, + }, + }, + 3: { + attention: { + attention: { + key: { + bias: Array([-0.22031489, 0.02730584, 0.23429762, -0.3125938 , 0.22592317, + -0.6771136 , -0.21247995, -0.341136 , 0.07302368, 0.44729188, + 0.22906882, -0.524169 , -0.04314182, -0.22182061, -0.67601067, + -0.01997758, -0.00890862, 0.31220517, 0.3243842 , -0.08772013, + 0.18925034, -0.43556792, -0.30325386, 0.13643748, 0.27908477, + 0.26076847, 0.19125256, 0.6231867 , 0.34878573, 0.21458477, + 0.12557186, 0.25881395, 0.11842591, 0.14228185, -0.14626004, + 0.02273282, 0.04946363, 0.3702873 , 0.4587399 , 0.46696767, + 0.2649066 , -0.02744056, 0.21641424, 0.4786166 , -0.39902607, + -0.25150967, 0.20756108, -0.9501807 , 0.2439707 , 0.07372095, + 0.1429905 , 0.5078466 , -0.47601166, 0.01509763, -0.63272023, + 0.5627352 , 0.5814683 , 0.3356227 , -0.07830408, -0.29685998, + 0.09222841, -0.49363494, -0.09527437, -0.01283222, 0.10691621, + 0.44450942, 0.2378995 , 0.14754085, 0.07311877, -0.4130656 , + 0.12185058, 0.1367209 , 0.18846467, -0.11077244, -0.10472129, + 0.13243866, -0.6801274 , 0.18685853, -0.40455377, -0.17242797, + -0.3273368 , 0.21752419, 0.5587714 , 0.33085388, -0.13658075, + 0.11016981, -0.14520068, -0.0551539 , -0.0988646 , 0.19754803, + 0.28224358, 0.27123928, 0.01996558, 0.25261068, 0.3291632 , + 0.73044914, -0.7135282 , -0.2712479 , 0.17239682, 0.4813283 , + 0.350012 , 0.76145405, 0.01346729, 0.28428316, 0.17260526, + -0.40829295, 0.66129416, -0.35163417, -0.15398747, -0.87850165, + 0.16401769, -0.36521718, 0.08581091, 0.12348412, -0.731215 , + 0.24665482, 0.77237046, 0.39620268, -0.4511514 , 0.34254473, + -0.669006 , -0.50109476, -0.18661116, 0.2477596 , 0.65949994, + 0.01988118, -0.04483053, -0.7621275 ], dtype=float32), + kernel: Array([[ 0.14074115, 0.11149215, 0.02833093, ..., 0.24423462, + -0.3480132 , -0.06901791], + [ 0.13614887, -0.2725882 , -0.31767422, ..., -0.30451745, + -0.0344476 , 0.2648898 ], + [ 0.19196813, 0.24347357, 0.09369045, ..., 0.3346163 , + -0.10084 , -0.19047955], + ..., + [ 0.02364591, 0.3210728 , 0.15941176, ..., -0.5476011 , + 0.30842453, 0.16627334], + [-0.3327589 , 0.2611727 , 0.04791668, ..., 0.5718048 , + -0.00553066, -0.12860055], + [ 0.13814813, -0.2723826 , -0.10484312, ..., 0.10856424, + -0.46137026, 0.3135762 ]], dtype=float32), + }, + query: { + bias: Array([ 0.01636963, 0.00303602, -0.2237812 , 0.17135 , -0.5504063 , + 0.71538967, 0.34468257, 0.6879721 , -0.32814807, 0.12084308, + -0.23467594, 0.24040897, -0.03144267, 0.41499555, 0.18407044, + -0.19885245, 0.29428306, -0.17422631, -0.72670996, 0.19728792, + -0.12558015, -0.04110402, 0.7530146 , -0.32672936, -0.34328216, + 0.39724907, -0.19968817, -0.0735909 , -0.11611477, 0.2887974 , + 0.27589712, -0.19556114, 0.00310108, 0.05990869, -0.06629719, + 0.11465285, 0.30233756, -0.4840145 , 0.36792257, 0.01198539, + 0.17746153, 0.00102492, -0.73527193, -0.04222427, 0.1278375 , + -0.04002868, -0.04392892, 0.07372226, 0.00746391, 0.01089494, + -0.06465401, -0.01894091, -0.04520243, -0.05135101, 0.10821903, + -0.36141765, 0.01786461, -0.00758973, 0.20370121, 0.07571532, + -0.04988701, -0.50066465, -0.29137757, 0.0468887 , -0.0369802 , + -0.28495568, -0.15369293, -0.04906723, -0.05060489, 0.23548573, + 0.02296117, 0.29195184, -0.4791245 , -0.32274112, 0.00163329, + 0.6024667 , -0.02721317, 0.4371868 , -0.45318347, 0.15298416, + 0.53022355, 0.23577152, -0.11320106, -0.05042947, 0.00792542, + 0.34630367, 0.45112005, -0.17176089, -0.30412972, -0.00977058, + 0.3797044 , -0.55427474, -0.14912651, -0.31753495, -0.14334434, + -0.12312431, 0.29336822, 0.08129798, -0.09826545, -0.37803188, + -0.03996738, -0.565557 , -0.15023306, -0.17836387, -0.00129342, + 0.31617662, -0.40511948, -0.0246349 , 0.8049514 , 0.50346047, + -0.08232877, 0.7958381 , 0.11525708, -0.87054974, 0.02044921, + 0.10254359, -0.41025582, -0.11836894, 0.3934775 , -0.2755682 , + 0.28480217, 0.28756127, -0.06242 , -0.36147273, -0.31789255, + -0.1169754 , 0.13862565, 0.5221684 ], dtype=float32), + kernel: Array([[-0.03366861, 0.1951855 , 0.13829306, ..., 0.20226802, + 0.14245163, 0.0851718 ], + [-0.12275442, 0.09568106, 0.1026165 , ..., -0.11756799, + -0.03399425, 0.1083973 ], + [-0.11897923, 0.44871688, 0.22821409, ..., -0.19063285, + 0.14112021, 0.05127604], + ..., + [ 0.42162004, -0.30832654, -0.31275368, ..., -0.48211238, + 0.221526 , 0.17857978], + [ 0.23314169, 0.38464716, 0.37320703, ..., 0.7464507 , + 0.65437895, 0.09728524], + [-0.02876819, -0.11770961, -0.18619719, ..., -0.4305184 , + -0.24435337, -0.0622173 ]], dtype=float32), + }, + value: { + bias: Array([ 3.80855761e-02, 1.19626271e-02, 9.54822153e-02, 3.39327864e-02, + 6.47817999e-02, 1.12312162e-04, 2.36410853e-02, 1.12216197e-01, + 1.71483278e-01, -4.48498651e-02, -6.95201615e-03, -5.79331703e-02, + 1.67439714e-01, -1.31781459e-01, 1.07828841e-01, -1.07748993e-01, + -2.16874084e-03, -2.65070498e-02, -1.59811676e-02, 7.09029958e-02, + 5.01474105e-02, -4.89775911e-02, 2.54220992e-01, -4.80323769e-02, + -1.79449767e-02, -5.27789593e-02, -4.24154028e-02, -1.70985550e-01, + -1.36494562e-01, 1.54837176e-01, -1.15875848e-01, -1.71242371e-01, + 5.86349219e-02, 1.91222817e-01, -2.35057212e-02, -8.43578875e-02, + -4.65269759e-02, -8.38292018e-03, -4.65860590e-02, -5.02779474e-03, + 4.50352840e-02, -2.31239274e-02, 5.93480729e-02, 5.71514890e-02, + 5.99062406e-02, 2.83261593e-02, -4.82074209e-02, -4.80451956e-02, + -6.26677051e-02, 5.32909296e-02, 6.67422032e-03, 4.72835749e-02, + 4.04058099e-02, -4.40718755e-02, -4.50864471e-02, 1.62461493e-02, + -5.04599810e-02, -3.60567644e-02, 4.91038598e-02, -7.88058713e-03, + 1.09715499e-01, 5.14943898e-02, 3.80876176e-02, 5.21656685e-02, + 1.56386979e-02, 6.83846846e-02, 7.60920197e-02, -8.33767503e-02, + -1.77296321e-03, -9.03388578e-03, -5.15193716e-02, 2.69258358e-02, + 4.77294379e-04, -2.96406243e-02, -8.59358441e-03, -2.96034720e-02, + -2.67447811e-02, -6.73482614e-03, 1.27622047e-02, -2.17594709e-02, + 8.59617069e-02, 1.03174895e-01, 2.63699442e-02, -1.10600488e-02, + -7.31578097e-02, -3.59972566e-02, 4.40209918e-02, -4.44998406e-02, + -1.85169447e-02, -4.11919132e-02, -3.58521268e-02, 1.47900637e-03, + 4.30759192e-02, 2.79228836e-02, 7.82085881e-02, 3.90527584e-02, + 3.11766323e-02, 2.85879765e-02, -8.20098892e-02, -7.59413168e-02, + -6.41580895e-02, 2.49105208e-02, -1.47475805e-02, -9.27800685e-02, + -1.73971094e-02, 8.51288959e-02, -2.54925918e-02, 4.36546579e-02, + -9.24658105e-02, -1.01948939e-02, -4.81284559e-02, -3.62936854e-02, + 2.90373899e-02, -4.27892022e-02, 3.08775119e-02, 5.98094938e-03, + 2.10134909e-02, 4.55508381e-02, 3.65342051e-02, 1.74732562e-02, + 1.06893666e-02, -1.42213493e-03, -2.61030179e-02, 2.22432707e-02, + 1.47231504e-01, 6.33382201e-02, 2.16640271e-02, 4.13991958e-02], dtype=float32), + kernel: Array([[ 0.00939619, 0.02482712, 0.01164271, ..., 0.0164959 , + 0.00769011, -0.03742562], + [ 0.18655218, -0.36816835, -0.08651444, ..., 0.2512268 , + 0.02853749, 0.00129432], + [ 0.09724977, -0.17311321, -0.15452595, ..., -0.4560767 , + -0.48346895, -0.4558498 ], + ..., + [-0.06453069, 0.4239059 , 0.45162615, ..., -0.27700382, + 0.13836603, -0.33547392], + [-0.02759907, -0.9605103 , -0.82787424, ..., -0.51210946, + 0.33198425, 1.0528204 ], + [-0.12098552, -0.10120461, -0.10430597, ..., 0.78033483, + -0.08541391, 0.05855659]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-1.72631651e-01, 2.81819075e-01, -1.31835248e-02, -6.50164206e-03, + -9.19895153e-03, -6.00125082e-02, 4.56261076e-02, -1.17319543e-03, + 3.18587162e-02, 2.30925724e-01, -2.78387107e-02, -7.64354365e-03, + -7.98858032e-02, 2.95766089e-02, -4.09963764e-02, 1.03628486e-01, + 5.56401908e-02, -8.70198682e-02, -1.29737146e-02, 1.11193627e-01, + -5.96015491e-02, -5.83102107e-02, -2.50512566e-02, -3.02468054e-02, + -3.29878554e-02, 4.26750369e-02, -3.15967947e-02, 1.79770100e-03, + 6.93126023e-02, 8.52573961e-02, 4.88307029e-02, -6.38186038e-02, + 1.12157442e-01, -9.84508917e-02, 5.63154779e-02, -7.03682825e-02, + 2.64801327e-02, -1.89797301e-02, 6.20345166e-03, -8.26654360e-02, + 2.39220913e-02, -9.58621129e-03, -8.39129508e-01, 4.56922092e-02, + -1.53480023e-01, 6.83549270e-02, 7.20679983e-02, 2.08752692e-01, + 4.47404385e-02, -5.50162159e-02, -5.08189276e-02, 5.22453226e-02, + -2.34786086e-02, 1.55734494e-01, -4.19281684e-02, -1.19545199e-02, + -8.40202495e-02, -8.96746367e-02, 6.41859993e-02, 7.71264508e-02, + 8.69383588e-02, -1.31521776e-01, -2.94522438e-02, 4.01933566e-02, + -1.60081685e-03, -3.63952629e-02, -5.23505732e-02, 8.07748511e-02, + 3.52987312e-02, 8.45032781e-02, -1.18315943e-01, 3.48625183e-02, + 5.90237789e-02, -7.91934412e-03, 8.07135582e-01, -6.57365983e-03, + -6.89911172e-02, 1.08262442e-01, -1.27490476e-01, 1.97968110e-02, + -6.29805326e-01, 5.13264053e-02, 3.17039192e-02, 1.63930492e-03, + 2.53316425e-02, 1.45042747e-01, 1.39836982e-01, 1.34169623e-01, + 4.38121334e-02, -1.00197628e-01, 3.51289548e-02, 3.51637453e-02, + -1.21273242e-01, -1.04602434e-01, -9.42274835e-03, -2.06037387e-02, + -4.44292687e-02, -2.39612963e-02, -1.44756138e-02, 9.55058485e-02, + 3.27826925e-02, 5.12430328e-04, 1.59999877e-02, -6.98894709e-02, + 1.01640105e-01, -5.48832715e-02, 6.11420302e-03, -7.91232064e-02, + -1.24512471e-01, -3.10322326e-02, -1.39527261e-01, 3.97435911e-02, + 1.18619666e-01, -1.12082511e-01, 1.25646042e-02, 1.72376752e-01, + 4.17521298e-02, -2.35856101e-02, -3.09479032e-02, 6.41909391e-02, + -2.64365948e-03, 2.64008511e-02, 1.21359736e-01, -4.78715599e-02, + 4.74941581e-02, 5.94909228e-02, -6.63625225e-02, -2.96208411e-02], dtype=float32), + kernel: Array([[ 0.30733457, 0.11394037, -0.11875194, ..., -0.02600583, + 0.26162902, -0.24450323], + [ 0.32579228, -0.01019196, 0.08372336, ..., -0.23464818, + 0.77629703, -0.14200607], + [ 0.40005246, -0.09073471, 0.01217228, ..., 0.12469898, + 0.4625953 , -0.2634586 ], + ..., + [-0.14610979, 0.26803997, -0.26479524, ..., -0.581134 , + 0.39575958, -1.1754091 ], + [-0.10739683, -0.4411798 , 0.68852794, ..., -0.01546267, + 0.06261325, -0.04542963], + [ 0.20287202, 0.13014646, 1.1704878 , ..., -0.26887265, + -0.26802245, 0.629474 ]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([ 1.71948429e-02, -1.09049575e-02, 6.19933335e-03, -3.97290401e-02, + -1.75763867e-04, 5.07760718e-02, 1.12305544e-02, -2.87154932e-02, + -2.46254611e-03, -1.75729260e-01, 9.95321851e-03, -1.72805071e-01, + -1.28897009e-02, 2.74409372e-02, -1.33280344e-02, 1.57673303e-02, + -2.70904624e-03, 1.06892176e-03, 2.89133284e-03, 1.36673786e-02, + 4.77137081e-02, -9.59468889e-06, 6.27165893e-03, -6.68704743e-03, + 1.20838732e-02, 1.39770750e-02, -3.87260579e-02, -5.88312745e-03, + 4.62619169e-03, 6.23012260e-02, 8.07429329e-02, 7.94423558e-03, + -2.50669499e-03, 1.25099311e-03, -6.49371138e-03, -9.59886145e-03, + 3.83810289e-02, 7.55729433e-03, 4.41730954e-03, 1.75167769e-02, + 3.34830210e-02, 9.28634033e-03, 1.21751567e-02, -1.52964927e-02, + -2.70186807e-04, 7.17550702e-03, 8.44905898e-02, -9.36958939e-02, + -1.57131581e-06, 4.74082539e-03, -8.86013955e-02, 7.37810042e-03, + -4.76047629e-03, 6.89025401e-05, -9.31238383e-03, 6.64415536e-03, + 7.76583795e-03, 2.29338557e-02, -1.50072891e-02, -3.97083163e-03, + 5.64871766e-02, 1.86805464e-02, 5.85068855e-03, 3.33538577e-02, + 3.15117603e-03, 3.10075725e-03, -3.19483913e-02, 3.20728831e-02, + 2.38816179e-02, 1.34205222e-02, -1.76607758e-01, -5.26308222e-03, + 2.82549043e-03, -1.20671205e-01, -1.07251152e-01, 3.50134517e-03, + 2.33329542e-04, -1.22928575e-01, -5.50287850e-02, -2.40997149e-04, + 4.36784141e-03, -3.80961671e-02, -1.24475360e-01, -1.74865853e-02, + 9.07492265e-02, -1.61794608e-03, 4.57073562e-03, -6.20833039e-03, + -5.41980378e-03, 2.80724547e-04, -4.64204960e-02, -4.63194735e-02, + 1.02222068e-02, -5.72193079e-02, -3.21766338e-03, 6.88658096e-03, + -4.31186426e-03, -3.94533984e-02, -1.81648252e-03, 5.71832731e-02, + -6.70271181e-03, 5.93083352e-03, -8.59189108e-02, -5.01573570e-02, + -3.59455287e-03, -8.11670348e-02, -3.07016950e-02, 2.01600092e-03, + 4.20039892e-02, -4.99186059e-03, -1.39956048e-03, -5.72945038e-03, + 1.23080192e-03, 1.21427193e-01, 1.27224373e-02, 5.32536209e-02, + -1.02717979e-02, -3.96808702e-03, 3.51909577e-04, 1.07945967e-02, + -4.73669693e-02, -5.23439702e-03, -3.12553300e-03, 6.25749454e-02, + -9.56553966e-03, -9.94754490e-03, 2.23868527e-02, -2.48279348e-02], dtype=float32), + scale: Array([ 4.24516723e-02, 1.21176187e-02, -4.84997500e-03, 5.57146311e-01, + 2.28432156e-02, 6.31359816e-01, -4.46775509e-03, 6.18052892e-02, + -4.33097221e-03, 2.61805564e-01, -1.88892316e-02, 6.95380867e-01, + -6.15348248e-03, 5.73870949e-02, 1.55512363e-01, -5.18870428e-02, + 1.60957292e-01, -1.62363220e-02, 4.54635024e-01, -1.21937275e-01, + 4.06006753e-01, 4.14683282e-01, 3.15475985e-02, 4.67635952e-02, + 1.88465670e-01, -1.32318744e-02, 4.83835369e-01, 8.42766091e-02, + -7.77654797e-02, 5.02031922e-01, 4.80292380e-01, 3.83823998e-02, + -4.09559421e-02, 1.13746740e-01, -1.26264850e-02, 7.38768047e-03, + 3.53148162e-01, -1.56045742e-02, 2.96747178e-01, 1.84800133e-01, + 2.40674987e-01, -5.23572266e-02, 1.17433036e-03, 6.38127849e-02, + 2.33978732e-03, -3.27690654e-02, 4.07037109e-01, 1.53259337e-01, + 3.08140907e-02, -5.92192961e-03, 3.31769854e-01, 3.03894043e-01, + -1.66736450e-02, -4.60295379e-02, 3.50147933e-01, 4.64291312e-02, + 7.68471509e-03, 9.69270691e-02, 1.85923763e-02, -9.66442749e-03, + 4.18877095e-01, 1.00185797e-01, 1.32984370e-01, 4.85280842e-01, + -7.05129327e-03, 4.97093098e-03, 3.28959554e-01, 9.48822647e-02, + 2.17689544e-01, 2.05316558e-01, 5.99317133e-01, 3.90748940e-02, + 5.26148966e-03, 3.09464663e-01, 3.66784409e-02, 4.52369526e-02, + -4.48510870e-02, 2.84545600e-01, 3.67504448e-01, 4.94225929e-03, + -5.27279684e-03, 8.71935904e-01, 4.51588362e-01, 2.58467104e-02, + 3.65057826e-01, 6.95686322e-03, -1.69470385e-02, 1.21857105e-02, + -1.65910535e-02, -8.54126003e-04, 5.98546147e-01, 3.60920489e-01, + 2.29799710e-02, 3.24622720e-01, -8.44317749e-02, -2.27606995e-03, + 3.23808223e-01, 2.70338178e-01, 5.69886751e-02, 1.97269484e-01, + 3.14833894e-02, -3.49436514e-02, 3.14999282e-01, 5.10905981e-01, + 2.21107565e-02, 3.05217803e-01, 2.38849521e-01, -1.36028053e-02, + 4.52475280e-01, 2.09674463e-02, -2.62735710e-02, -6.97454140e-02, + -8.80253036e-03, 4.10522908e-01, -2.87996326e-03, 2.52809376e-01, + -1.01799294e-02, 1.55074904e-02, 4.01195139e-02, 1.15418859e-01, + 1.29438698e-01, -9.92912054e-03, 3.01613122e-01, 9.44039226e-02, + -3.23632099e-02, 1.09247090e-02, 5.09600282e-01, 6.76498190e-02], dtype=float32), + }, + layernorm_before: { + bias: Array([ 2.0053135e-02, -2.0385702e-01, -5.2851941e-02, -1.5826782e-02, + -3.9745602e-03, -5.1497813e-02, 4.4322563e-03, -2.0497143e-02, + 3.5912704e-02, -1.5449695e-01, -6.0098185e-03, 4.4156699e-03, + -5.7948790e-03, 1.7218473e-03, 8.6186066e-02, -8.2310978e-03, + 8.4572081e-03, 3.2470144e-02, -7.2876806e-04, -7.4659437e-02, + -1.4426826e-02, -6.7664213e-03, 1.5026542e-02, 4.0727290e-03, + 2.3569600e-04, -9.6417527e-04, -1.1503830e-03, -1.6481975e-02, + -7.0209317e-03, 5.3436416e-03, 6.8159692e-02, -2.5559873e-03, + -8.0474456e-05, -3.0456469e-03, -3.3534174e-03, 6.1106024e-04, + -6.6513027e-04, -1.2051529e-01, -4.2777561e-02, 6.5030590e-02, + 2.1362038e-02, 3.2946058e-02, 2.9181463e-01, 1.6087948e-03, + 8.0183737e-02, 6.0054217e-04, -1.5031694e-03, -1.3727798e-01, + 2.0673402e-02, 1.6634881e-02, -6.1021522e-02, -9.7963065e-03, + -4.4847801e-03, -1.6810349e-03, -6.0316608e-03, -1.3532469e-04, + 8.6372629e-02, -1.0683674e-02, -1.5961612e-02, 4.4840686e-02, + 5.4682657e-02, 1.3794379e-02, 2.7690900e-03, 3.4621611e-02, + 3.8512382e-03, 1.3362155e-03, 2.2860983e-02, 1.7252501e-02, + 7.1021970e-03, -1.8353600e-02, -9.1281412e-03, -2.5052501e-02, + 6.4535923e-02, -1.1188607e-01, -7.4342960e-01, 9.8237805e-03, + -4.8992522e-02, -3.0361801e-02, 4.7444657e-05, -1.2857079e-02, + -3.5183441e-03, -2.8601723e-02, -4.8774560e-03, -4.9416926e-02, + 2.1382656e-03, -1.7278522e-02, 4.1791960e-03, -6.8051869e-04, + 5.9474696e-04, 1.3638437e-01, -9.9423109e-03, 1.4076917e-03, + -2.9004067e-03, 1.2549493e-02, 1.8330464e-02, 9.2371933e-02, + -8.4709952e-04, -1.5642725e-02, 1.8227950e-02, 4.0236771e-02, + -4.6799608e-02, -6.4027351e-03, -1.0166033e-01, 2.3036268e-02, + -2.3710703e-02, -2.6576657e-02, 4.9998323e-03, 2.3283776e-02, + 6.9294289e-02, -2.0691980e-02, 4.1049831e-02, 1.8761930e-04, + -1.0511643e-01, 8.2109414e-02, -1.6860457e-02, -3.8362190e-02, + 4.8856564e-02, -2.3737226e-02, -1.6952256e-02, 1.5742332e-02, + -1.9643296e-02, 7.6947906e-03, 1.0548260e-03, 5.7553727e-02, + 1.7045023e-02, -3.2199997e-02, 7.6176226e-02, -7.1644776e-02], dtype=float32), + scale: Array([ 1.36227664e-02, 1.40898958e-01, 1.92645520e-01, 4.70757820e-02, + -2.37463228e-03, 3.57880324e-01, 1.77441984e-01, 2.04336628e-01, + 5.04735708e-02, 2.16793969e-01, 1.84319422e-01, 1.62005424e-01, + -2.82178313e-04, 3.39963939e-03, 3.45121980e-01, 3.06478173e-01, + 4.18081768e-02, 7.82519281e-02, -1.85935143e-02, 2.79815465e-01, + 2.39701569e-01, 2.40149871e-02, 9.43208337e-02, 1.05237709e-02, + 9.96645987e-02, 5.39937872e-04, -1.31208654e-02, 1.39623404e-01, + -3.63838188e-02, 3.99398878e-02, 3.38648826e-01, 9.40667465e-02, + 1.93201587e-01, -4.52386551e-02, 3.65994535e-02, 3.85606177e-02, + -3.78566347e-02, 2.66506523e-01, 2.40695551e-01, 3.16709459e-01, + 8.35959762e-02, 2.61536777e-01, 1.21971466e-01, 1.21512767e-02, + 3.52942318e-01, 2.85213944e-02, -1.15962522e-02, 2.87642092e-01, + 2.67458588e-01, 3.02993476e-01, 1.73773259e-01, 1.12724066e-01, + -8.12414661e-03, 2.45114770e-02, 1.12504624e-01, 2.88635880e-01, + 1.51033133e-01, -3.74849811e-02, 3.15434486e-01, 3.69797438e-01, + 3.93722087e-01, 1.04381077e-01, 1.80823863e-01, 2.28892982e-01, + 2.66926765e-01, -5.73799293e-03, 1.90365806e-01, 2.51159281e-01, + 2.34797284e-01, 2.54733443e-01, 1.83505401e-01, 1.44614145e-01, + 2.63359755e-01, 4.59246188e-01, 2.01807499e-01, 2.74832934e-01, + 2.06239209e-01, 1.09179243e-01, 6.59588631e-03, 2.21557349e-01, + -5.72106568e-03, 2.62187541e-01, 5.55602312e-02, 2.09347799e-01, + 6.81416085e-03, 6.24816194e-02, -1.57044847e-02, 5.38076796e-02, + -1.21786725e-03, 3.13683510e-01, 7.88503736e-02, 1.91961545e-02, + -7.89848214e-04, 2.04859436e-01, 1.77144408e-01, 1.27074376e-01, + -2.01367047e-02, 1.25420019e-01, 6.45905286e-02, 2.66942471e-01, + 3.49235058e-01, 2.52543539e-01, 3.24409813e-01, 2.59198427e-01, + 3.92820060e-01, 1.41132906e-01, 9.48882196e-03, 1.74179897e-01, + 5.39034963e-01, 1.58920869e-01, 2.73100972e-01, 1.81179568e-02, + 1.68145299e-01, 3.43630970e-01, 1.51725076e-02, 2.98200041e-01, + 2.57570058e-01, 2.81464636e-01, 2.47943982e-01, 1.76195994e-01, + 5.23593724e-02, 1.48531556e-01, 2.94840173e-03, 2.51852840e-01, + -7.62507394e-02, 1.64415836e-01, 4.18539464e-01, 1.66102067e-01], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-8.55890103e-03, 3.67313088e-03, -4.36119083e-03, -8.68369080e-03, + -1.34754460e-03, -6.52316585e-03, -5.01648569e-03, -8.97654798e-03, + 3.13179940e-03, -7.83655513e-03, 2.37924187e-03, -5.66111412e-03, + -1.27145583e-02, -8.22409242e-03, 2.74701975e-03, -4.37194854e-03, + -4.04660078e-03, -3.13023920e-03, -1.81288610e-03, -9.03141592e-03, + -1.66141614e-03, -5.17393708e-01, -1.02338977e-02, 3.56379431e-04, + -7.41085410e-03, 9.37872624e-04, -2.64340243e-03, 5.81945898e-03, + -1.19625553e-02, -4.65264637e-03, -6.52810233e-03, -7.84709677e-03, + 1.05619885e-03, -9.87111218e-03, 4.54840425e-04, -5.49622951e-03, + 7.56876124e-03, -5.51949022e-03, -3.83866858e-03, -1.84899215e-02, + 3.52073257e-04, -2.78544938e-03, -5.13466494e-03, -4.08213615e-01, + -5.55376470e-01, -1.75705249e-03, 4.89682052e-03, 1.53476396e-03, + 1.07086054e-03, -3.85978585e-03, 9.40655766e-04, 6.59122784e-03, + -3.99606815e-03, -6.94369990e-03, -2.59699044e-03, 7.18095253e-05, + -8.83713830e-03, -1.96340843e-03, -2.48700310e-03, -9.87609662e-03, + -6.28701295e-04, -7.88388960e-03, 7.09942728e-03, 9.48800612e-03, + -1.95277121e-03, -1.22706648e-02, 2.85065360e-03, 9.00720805e-03, + -7.02000320e-01, 4.25555557e-03, -1.28851901e-03, 2.98675569e-03, + -3.15172016e-03, -1.08107440e-02, -8.73899739e-03, -1.31636495e-02, + -9.91051551e-04, -1.34902084e-02, 3.62220523e-03, 3.06941615e-03, + -1.32400328e-02, -2.95568615e-01, 7.01408414e-03, -5.55180886e-04, + 3.73574020e-03, -7.83216953e-03, 7.60261260e-04, -9.49840934e-04, + -1.34194596e-02, -1.67394914e-02, -8.75537097e-03, -1.25364307e-03, + -6.53209072e-03, -1.00438660e-02, -4.49657347e-03, -1.30805969e-02, + 7.16269715e-03, -1.18865566e-02, 1.15038871e-04, -7.90204573e-03, + -1.13492664e-02, -1.28111634e-02, -2.23936560e-03, -5.98883303e-03, + -1.25457710e-02, -1.07937502e-02, -3.34314862e-03, -7.55358534e-03, + -6.86253095e-03, -2.35536560e-01, -9.40762553e-03, 1.64968718e-03, + -2.74949777e-03, -3.80960264e-05, 1.37860198e-02, 1.36729311e-02, + -1.17542176e-03, -4.75597620e-01, 1.06082670e-02, -4.49829688e-03, + 1.80422806e-03, -3.01245158e-03, -5.77643573e-01, -8.06397665e-03, + 1.37929013e-03, -9.44454595e-03, -2.07981328e-03, 9.04347980e-05, + -5.42783411e-03, 1.13049364e-02, -6.55044988e-03, -9.25507303e-03, + -1.70839131e-02, -3.13228741e-03, -1.04950890e-02, -9.20292269e-03, + -4.24578832e-03, -8.48777406e-03, -1.27782347e-02, -7.17586139e-03, + -5.02940093e-04, -6.85630506e-03, -5.49637014e-03, 5.60814887e-03, + -7.42090074e-03, -1.33384299e-02, -4.22688900e-03, -3.86846159e-03, + -1.31040469e-01, -5.72234035e-01, 1.44373742e-03, -2.45060981e-03, + -5.27091743e-03, -1.61155630e-02, -8.04609526e-03, -3.71574331e-03, + 4.77314461e-03, -2.98696361e-03, -1.80231046e-03, -3.76312784e-03, + 2.70129368e-03, -5.70825338e-01, -1.65502690e-02, -6.61239505e-01, + -9.16241575e-03, -1.41646145e-02, -5.00553148e-03, -4.18055616e-03, + -4.17894637e-03, -1.43527528e-02, 1.20390498e-03, 1.38255325e-03, + -2.05959729e-03, -1.65631052e-03, 8.46141134e-04, -4.89455275e-03, + -3.43874632e-03, -4.29662468e-04, -9.45413113e-03, -4.12517926e-04, + -2.43844930e-03, -8.26839264e-03, -6.06988370e-03, -1.55913727e-02, + -6.00419287e-03, 7.41206226e-04, -4.03943751e-03, -8.32396373e-03, + -1.77192141e-03, -8.27328023e-03, -4.80760727e-03, -1.19567281e-02, + -3.16868187e-04, -1.15141738e-02, 7.74921337e-03, -2.20119534e-03, + 1.03418762e-03, -5.75479446e-03, -4.55250964e-03, -1.04654999e-02, + -6.14567241e-03, -1.54978642e-02, -7.39947148e-03, -1.07835466e-02, + -1.21487593e-02, -6.76209899e-03, 1.25275925e-03, -3.19723099e-01, + -3.41119803e-03, 2.97757820e-03, -5.52376360e-03, -9.89249256e-03, + -4.31803286e-01, -1.36245508e-02, -7.62469275e-03, -1.74430560e-03, + -1.96244684e-03, 5.65344701e-04, -6.50570095e-01, -1.45257404e-03, + -2.76738009e-03, -8.20099376e-03, 2.03965130e-04, -1.13646872e-02, + -2.37802067e-03, -5.39474888e-03, 2.26747943e-04, -2.44784285e-03, + -1.09350355e-02, 8.08892888e-04, -8.32007360e-03, -1.91815558e-03, + -9.69145447e-03, 1.22545648e-03, 4.08804649e-03, -5.84384939e-03, + 3.66018945e-03, 4.44652513e-03, -9.09210462e-03, -4.17519035e-03, + -3.98625523e-01, -3.37073533e-03, -8.35657492e-03, -4.79207002e-03, + -1.39846019e-02, -8.26936215e-03, -8.90700985e-03, -2.52558112e-01, + 3.46473581e-03, -1.71433843e-03, -1.91845442e-03, -3.26462858e-03, + -8.05919431e-03, -5.96231222e-03, -2.60497467e-03, -5.12533821e-03, + -9.94778052e-03, -1.35790939e-02, -7.90751539e-03, -5.66528225e-03, + 2.63936934e-03, -4.14405525e-01, 2.22919116e-04, 8.12191633e-04, + -1.82921547e-04, -7.25358841e-04, -1.08812936e-02, 3.49700474e-03, + -1.17458524e-02, -1.99729693e-03, -5.74411824e-03, -1.04112234e-02, + -5.87427756e-03, 5.42524364e-03, -4.66417568e-03, -3.61002120e-03, + -4.92561430e-01, -5.91339357e-03, -8.21670890e-03, 6.48486370e-04, + -1.16507420e-02, -7.28300912e-03, 9.69227776e-03, -1.40586449e-03, + 1.34252605e-03, -4.55454225e-03, -1.05377687e-02, 8.61465523e-04, + -1.29449442e-02, 3.23367422e-03, -1.80069392e-03, 2.35577201e-04, + -8.91069509e-03, -1.41903665e-03, -8.80321302e-03, -4.76717623e-03, + -2.86917109e-03, -3.21599003e-03, 3.61907203e-03, -7.99222197e-03, + -6.45599188e-03, -5.55047020e-03, -9.93613340e-03, -8.86950456e-03, + 1.25965700e-02, 3.33399116e-03, -1.01383878e-02, 2.34254939e-03, + 5.69945900e-03, -5.75918052e-03, 7.65805482e-04, -1.72046013e-02, + -8.67277035e-04, -1.03675043e-02, -7.15132989e-03, -1.15488358e-02, + -2.48059118e-03, -1.51636070e-02, -1.34787019e-02, -8.78769148e-04, + -6.11743808e-01, -6.02744520e-04, 7.56409660e-04, -3.29326536e-03, + 5.58212679e-03, -9.33979824e-03, -2.45822361e-03, -5.48888650e-03, + -4.64835018e-03, -4.38324921e-03, -3.24640493e-03, -6.00662351e-01, + 1.93342322e-03, -1.04327723e-02, 4.34050598e-04, 3.05471686e-03, + -7.74156442e-03, -1.68287978e-02, -1.33079069e-04, -3.91604519e-03, + -1.49135012e-02, -8.73885211e-03, -5.66676259e-03, 3.75113130e-04, + -5.36525948e-03, -1.39129013e-02, 1.88402925e-03, -6.80929542e-01, + 2.44663190e-03, -1.16351806e-02, -5.38230175e-03, -3.69915483e-03, + -4.62271646e-03, 4.26167302e-04, 1.99855946e-04, -1.86087582e-02, + 7.48794526e-03, -1.61632244e-02, -4.98278052e-01, -8.40996485e-03, + -6.01133797e-03, -8.28853925e-04, -5.63414209e-03, 5.07016899e-03, + -9.90651455e-03, -4.54475218e-03, 1.22021348e-03, 5.08222263e-03, + -5.74093172e-03, 2.58398801e-03, -5.80827938e-03, 2.76929396e-03, + -8.52370542e-03, -8.13993160e-03, -3.10428627e-03, 9.73783247e-03, + -2.26997398e-03, -7.61909829e-03, 4.68078069e-03, -3.60694993e-03, + -8.70567188e-03, -1.44884512e-02, -2.72786571e-03, 4.09194687e-03, + -9.40709934e-03, 4.71121492e-03, -5.04164118e-03, -1.17910588e-02, + 1.36497710e-03, -1.07826991e-02, -1.53774535e-03, -1.12592215e-02, + -2.51332531e-03, -1.03800860e-03, 9.40567255e-03, -1.59978624e-02, + -9.16326977e-03, 5.42400870e-04, -2.22451333e-03, -7.41019659e-03, + -3.90202254e-01, -1.51460301e-02, -6.29314408e-03, -4.74618049e-03, + 3.59613076e-03, -1.17005240e-02, -1.27715422e-02, -1.15625896e-02, + -1.68063808e-02, -1.12238340e-02, -5.49799949e-03, 7.02496106e-03, + -2.54716747e-03, -9.75850690e-03, -3.14950384e-03, -7.65251042e-03, + -1.46619789e-03, 1.40559685e-03, -2.16309223e-02, 4.46918281e-03, + -8.65491666e-03, 5.14270552e-03, 3.51362582e-03, -9.26864985e-03, + -1.17338272e-02, -2.89736595e-03, -1.97741878e-03, -5.52454265e-03, + -8.56413040e-03, -9.34949901e-04, 2.97652930e-03, -2.00355076e-03, + -3.76361725e-03, 1.10566849e-03, -5.63339936e-03, -2.81627057e-03, + -1.68547058e-03, -4.57938435e-03, 6.58686552e-03, -4.55372721e-01, + -5.61221363e-03, -6.51222235e-03, 1.07192027e-05, -8.45061056e-03, + 5.90987783e-03, -5.17501831e-01, -4.53207910e-01, 1.17717020e-03, + -7.01933773e-03, 4.97156382e-03, 1.22358715e-02, -3.75395943e-03, + -1.19871208e-02, 7.76410243e-03, -1.40578160e-02, 3.29062645e-03, + 6.30691764e-04, -1.67252067e-02, 1.95841817e-03, -9.67236888e-03, + -1.17397988e-02, -6.99722324e-04, -6.32376294e-04, -6.49616960e-03, + -2.94766188e-01, 1.18942198e-03, -7.66102644e-03, 3.43931909e-03, + 2.72910809e-03, -5.50179362e-01, -1.15189431e-02, -5.44976583e-03, + -3.17149970e-05, 1.41578354e-03, -6.40626810e-03, -3.47646652e-03, + 6.37558231e-04, -2.99780950e-04, 5.98926889e-03, -4.70549008e-03, + -7.91283231e-03, -9.49787791e-04, -3.55603476e-03, -1.14685595e-02, + -1.00854738e-03, 2.23665382e-03, -2.58625904e-03, -1.55551946e-02, + -9.62937716e-03, -3.64469341e-03, -6.78142011e-01, -3.52118816e-03, + -1.36016486e-02, -1.84007958e-01, -5.86909475e-04, -8.81857052e-03, + -1.06971972e-02, -7.59707938e-04, -5.68926990e-01, -5.22532880e-01, + -2.09414326e-02, -1.72124021e-02, -6.04376057e-03, -3.09983082e-03, + -4.42491425e-03, -8.42403620e-03, 7.68742245e-03, -3.21401260e-03, + 6.66661141e-03, 2.13651289e-03, -4.34801867e-03, -7.12090218e-03], dtype=float32), + kernel: Array([[-0.06747823, 0.12502806, 0.01688226, ..., 0.04402845, + 0.0489017 , 0.02632801], + [-0.01445617, 0.03798477, -0.00418379, ..., 0.01224544, + 0.02048267, 0.01027028], + [-0.01829829, 0.01891246, 0.01610937, ..., 0.03225997, + -0.01639963, -0.03890606], + ..., + [ 0.00887221, -0.00487363, 0.02399558, ..., 0.01346108, + 0.01049829, -0.03240183], + [-0.01674049, -0.00387911, -0.03074075, ..., 0.00146568, + -0.00876305, -0.02836148], + [ 0.01625663, 0.03327096, -0.04843638, ..., 0.08134065, + 0.03410971, 0.01957767]], dtype=float32), + }, + output: { + bias: Array([-0.16341606, 0.2722003 , -0.01811248, 0.01928756, -0.02578784, + -0.10320964, 0.0457683 , 0.00972074, 0.05574008, 0.24268645, + -0.02999464, 0.043408 , -0.0737564 , 0.03339861, -0.04492559, + 0.09947823, 0.05016089, -0.08309133, -0.05127812, 0.11153366, + -0.08190622, -0.06840179, -0.02421227, -0.02895692, -0.04004822, + 0.03941108, -0.0170032 , -0.00097349, 0.07023378, 0.02564083, + 0.03406862, -0.06778747, 0.11675587, -0.08771975, 0.04982556, + -0.07175723, -0.04058652, -0.01989682, -0.02999596, -0.07330333, + 0.03119163, -0.0147863 , -0.83884764, 0.04124971, -0.147959 , + 0.06891607, 0.05300924, 0.2076162 , 0.04646805, -0.05764642, + -0.03448724, 0.05170846, -0.02983713, 0.14581041, -0.04477751, + -0.0028539 , -0.08908532, -0.09743858, 0.05685575, 0.08213867, + 0.0757231 , -0.13601153, -0.02382083, -0.01865633, -0.00175037, + -0.03086867, -0.03948168, 0.06656119, 0.03003574, 0.07914743, + -0.02571616, 0.03044318, 0.05824798, 0.02247159, 0.7948062 , + -0.00407046, -0.07054033, 0.12808433, -0.08461454, 0.02311847, + -0.65620077, 0.10964675, 0.10405939, -0.00352513, -0.02670519, + 0.15050422, 0.13908558, 0.14308207, 0.03576402, -0.11146745, + 0.07615335, 0.04982926, -0.11778215, -0.10201599, -0.01381564, + -0.01693749, -0.03509882, 0.01073222, -0.01447865, 0.08480705, + 0.02910439, 0.00359502, 0.01869445, -0.05797101, 0.09884419, + 0.0161234 , 0.00668302, -0.08439763, -0.13141344, -0.03596077, + -0.14647698, 0.03868678, 0.11478355, -0.14481665, 0.00641343, + 0.16066928, 0.04811645, -0.02036389, -0.03384693, 0.06242058, + -0.00735435, 0.02676932, 0.12175596, -0.06479023, 0.04808705, + 0.05643554, -0.07896087, -0.02733074], dtype=float32), + kernel: Array([[ 0.04431782, -0.07884661, -0.02030264, ..., 0.03209 , + 0.02117256, 0.00271614], + [-0.05527279, 0.03204606, -0.00293706, ..., -0.03798404, + -0.02230972, 0.00913886], + [-0.01904109, -0.02136857, 0.00946268, ..., 0.00186177, + -0.02516458, 0.05516352], + ..., + [-0.03168841, 0.07299381, -0.00359565, ..., -0.01263119, + -0.04168412, 0.01372963], + [-0.02394411, 0.08588913, 0.00119561, ..., 0.0126469 , + -0.04365737, 0.0029881 ], + [-0.01780676, 0.05708533, 0.01049639, ..., 0.03207007, + -0.03454534, 0.0507676 ]], dtype=float32), + }, + }, + }, + }, + 4: { + attention: { + attention: { + key: { + bias: Array([ 0.25172058, 0.68627775, 0.5949247 , 0.1658349 , -0.3173467 , + 0.1774752 , 0.25412992, 0.08010545, 0.56002 , 0.17958656, + -0.17482378, 0.07066374, 0.53656256, -0.5385957 , 0.03753721, + 0.3764516 , -0.05485516, 0.7278333 , -0.7301842 , 0.7662722 , + 0.8557523 , 0.19676152, -0.7459649 , 0.11304809, -0.27465415, + 0.41445032, 0.31016302, -0.9027352 , 0.15300739, -0.12067419, + 0.11333649, 0.36544785, 0.10615282, -0.19266807, 0.00742444, + -0.3133698 , 0.13811715, 0.17462967, -0.04043375, 0.01431966, + 0.02120961, -0.31457424, 0.24998754, -0.30905357, -0.33957723, + 0.44098425, 0.31071022, 0.5162943 , 0.1519571 , 0.20767058, + 0.11702721, 0.30976802, -0.26687843, 0.02610935, -0.38046166, + 0.15602328, 0.1049099 , -0.3816162 , 0.34945863, -0.13488789, + 0.25807512, -0.14170453, -0.2737117 , 0.6794868 , -0.5637404 , + -0.7506761 , -0.59640557, -0.35966128, -0.36779413, 0.6445572 , + -0.15863825, 0.28188404, 0.6817374 , -0.30411932, 0.15068318, + 0.21245846, 0.47254717, -0.3180801 , -0.72213024, -0.00796549, + 0.52852696, -0.74865776, 0.72052324, -0.6231846 , -0.35524404, + -0.6322673 , 0.5927172 , 0.7579426 , 0.22011517, 0.37161058, + -0.58662343, -0.5653432 , 0.4952815 , 0.34009102, 0.4199334 , + -0.04050065, 0.967862 , 0.62645787, -0.24915175, 0.6059676 , + -0.12386217, -0.97235304, 0.98436683, -0.71207863, -0.2554404 , + 0.29857978, -0.6535443 , -0.00581208, 0.54410416, -0.64588374, + -1.2313039 , -0.7612568 , -0.5072008 , -0.15740068, 0.5934992 , + 0.29966694, -0.78861797, 0.95648044, 0.90275365, 0.3240316 , + -0.7827806 , -0.90508664, 0.24878727, -0.92607623, 0.45041922, + -0.7728655 , 0.2784089 , 0.75287944], dtype=float32), + kernel: Array([[ 1.6211687e-01, -2.6181787e-02, -5.1868878e-02, ..., + -2.0115136e-01, -1.6275570e-01, 4.0626042e-02], + [ 2.0827402e-01, -1.1491413e-01, -1.4099246e-01, ..., + 1.2704584e-01, 2.0065814e-01, 3.5636088e-01], + [ 1.4854251e-01, -1.1401314e-01, -4.7746781e-02, ..., + -8.4292240e-02, 1.9939394e-01, 9.9089958e-02], + ..., + [-5.0814073e-03, 7.1649961e-02, 8.1174605e-02, ..., + 1.8439499e-01, -1.8450679e-01, 3.2407336e-04], + [ 5.2984446e-02, -4.0257105e-01, -3.7533784e-01, ..., + 2.7334958e-01, -7.9834215e-02, -2.6604572e-01], + [-1.3808161e-01, -2.5356847e-01, -2.2354984e-01, ..., + -3.9307380e-01, 3.9180964e-01, 4.5249498e-01]], dtype=float32), + }, + query: { + bias: Array([ 3.67379338e-01, -3.86799455e-01, -3.21791142e-01, -8.42057094e-02, + 4.18459654e-01, 5.85961580e-01, -4.60834682e-01, -6.04065180e-01, + 1.94889054e-01, 5.26792884e-01, -4.81182262e-02, -8.55236873e-02, + -5.21454439e-02, 4.12250131e-01, -1.91803545e-01, -5.08220792e-01, + -6.88099623e-01, 1.73875868e-01, -6.42500934e-04, 3.13504860e-02, + 7.57871941e-02, -1.18366331e-01, 7.69717107e-03, 6.50874972e-01, + -1.46498501e-01, -3.85033339e-02, 1.68164968e-01, 2.94075936e-01, + 5.43904722e-01, 1.82771876e-01, 4.28779982e-02, 3.94661844e-01, + 2.40447834e-01, 7.63046503e-01, 1.41106710e-01, 1.50562599e-01, + 7.20623806e-02, -3.31430919e-02, 9.66051295e-02, -4.25018109e-02, + 2.59480290e-02, 5.50236627e-02, 6.59183860e-02, 7.35617638e-01, + 9.20522138e-02, -1.12395249e-01, -2.63706028e-01, -4.15704325e-02, + 9.31539387e-02, -6.19517326e-01, 3.92871583e-03, -2.88900491e-02, + 8.03363383e-01, -1.52911708e-01, 7.79530257e-02, -1.06074944e-01, + -5.33656120e-01, 1.22914307e-01, 1.09775752e-01, 3.54613900e-01, + -4.29570854e-01, 1.18217006e-01, 5.40162444e-01, -1.17120035e-01, + -6.53074026e-01, -4.32521045e-01, -7.02763855e-01, -4.33477342e-01, + -4.24702972e-01, 4.39991415e-01, -2.19614744e-01, 4.09397632e-01, + 5.78722954e-01, 1.95280928e-02, 2.52549767e-01, 5.14830410e-01, + 5.46005607e-01, -6.23379529e-01, -4.24228251e-01, -2.12890625e-01, + 2.95019180e-01, -6.27236724e-01, 5.73115885e-01, -3.63081634e-01, + -3.95359695e-01, -4.50571030e-01, 5.81661999e-01, 5.84261417e-01, + 2.53536344e-01, 5.42065144e-01, -3.84602755e-01, -4.97589141e-01, + 4.21692282e-01, 4.60142285e-01, 3.82059246e-01, -3.72222006e-01, + 5.13080180e-01, 2.00476050e-01, -6.01216495e-01, 8.54436874e-01, + 2.28271738e-01, -3.90953124e-01, -9.38059986e-02, -8.37827742e-01, + 5.72855584e-02, -1.37482524e-01, 1.96730390e-01, 1.36508822e-01, + 1.39595941e-01, -4.52210784e-01, 1.51813626e-01, -7.74807990e-01, + 2.37015381e-01, -3.03821355e-01, -3.49135190e-01, -1.64326906e-01, + -8.14965189e-01, -1.61662132e-01, 8.50251555e-01, 2.36134544e-01, + 1.86818361e-01, -8.37599099e-01, 9.94282067e-02, 8.31496865e-02, + 1.03579015e-01, -4.60613877e-01, 8.01612064e-02, 8.32802534e-01], dtype=float32), + kernel: Array([[ 0.00486563, 0.01157592, 0.0173437 , ..., 0.06944456, + -0.07013498, 0.01022998], + [-0.07428966, -0.05709548, -0.03850973, ..., -0.09689807, + 0.22320649, 0.2770357 ], + [-0.10973258, 0.19551165, 0.12421291, ..., -0.02782791, + -0.36778155, -0.19087008], + ..., + [ 0.30281523, 0.3621789 , 0.32462892, ..., 0.0546729 , + -0.04859495, -0.10934605], + [ 0.5838381 , 0.22978936, 0.17813566, ..., -0.16678724, + 0.70342004, 0.18783332], + [-0.48419687, 0.04412645, -0.01096392, ..., -0.05038855, + 0.37432894, 0.02556431]], dtype=float32), + }, + value: { + bias: Array([ 0.00544928, 0.17689343, 0.04735808, -0.18264541, -0.04237033, + -0.17941093, 0.00340467, -0.00149406, -0.23728922, -0.07449124, + -0.2651422 , -0.0984367 , -0.12018497, -0.08068941, -0.01734875, + -0.06606279, 0.1477764 , -0.02022878, -0.1402983 , -0.02072977, + 0.12049822, 0.06401265, 0.00406305, 0.15095161, -0.20597391, + 0.21405579, -0.08197237, -0.00951357, -0.13306661, -0.08558191, + 0.08653639, -0.10818963, -0.1644184 , -0.15369424, 0.3490528 , + 0.01967676, -0.16458155, 0.10356332, -0.2320539 , -0.0433312 , + 0.294788 , 0.08036903, -0.12358049, 0.16056572, 0.25535643, + 0.16914676, -0.06201077, 0.28501964, -0.13111252, 0.11861207, + 0.26627177, 0.16901444, -0.19237208, -0.08137509, 0.2016543 , + 0.15213989, -0.14299503, -0.24307302, -0.29300013, -0.0722383 , + 0.05268372, -0.14988767, -0.19426815, -0.24243473, -0.06109034, + -0.13096058, 0.155086 , 0.15769799, 0.10678913, 0.00772086, + 0.12339002, -0.02504941, -0.13470462, 0.02844545, 0.09698638, + 0.11937644, -0.01387584, -0.13328613, -0.05046262, -0.02981377, + -0.11246655, -0.06740061, -0.00562667, 0.15140204, 0.08057123, + -0.03175431, 0.05622792, 0.11395832, -0.09259764, -0.02921034, + -0.10965273, 0.09367548, 0.07572406, 0.17622064, -0.12645517, + -0.06927612, -0.13641733, 0.11861874, -0.01573502, -0.10696157, + -0.14588912, -0.16159363, 0.1815869 , 0.08687254, 0.11738776, + -0.0497781 , -0.12963611, 0.01230228, -0.10121961, -0.06638794, + 0.11086646, -0.05197389, 0.24918154, -0.16594617, 0.20172946, + 0.06618202, 0.1124498 , 0.11753738, -0.19368784, 0.04725859, + -0.13396324, -0.10240424, -0.11015446, 0.0015793 , 0.05593257, + -0.12091528, 0.08938907, -0.03860855], dtype=float32), + kernel: Array([[ 0.06369832, 0.04487664, -0.0572863 , ..., -0.00566542, + 0.00958204, -0.0823305 ], + [-0.48305708, -0.05723378, 0.09843583, ..., 0.01919795, + 0.04487755, 0.36410812], + [-0.39578405, 0.17672975, -0.04119537, ..., 0.06564505, + -0.01046039, -0.30802208], + ..., + [-0.04638171, -0.03762753, -0.15475406, ..., 0.11968867, + -0.1331625 , 0.01799979], + [ 0.34840006, 0.0131262 , -0.17438433, ..., 0.21620421, + -0.14971787, -0.73045504], + [ 0.14833876, 0.18964465, 0.05431836, ..., 0.10322906, + 0.10638059, -0.31155896]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-1.46490365e-01, 2.10801542e-01, -1.72206163e-02, -3.77512947e-02, + -7.96984136e-03, -7.52876028e-02, 1.01055026e-01, 2.95142401e-02, + 4.51585725e-02, 2.89253950e-01, -3.78346108e-02, 3.64244655e-02, + -6.14614636e-02, 4.95899320e-02, -3.73293981e-02, 5.79853542e-02, + 8.40439461e-03, -1.02269538e-01, -7.46445209e-02, 1.14303716e-01, + -2.32893787e-02, -3.25526744e-02, -1.71621908e-02, -2.69102436e-02, + -4.24926169e-02, 6.34983480e-02, -9.30675417e-02, -2.39975707e-04, + 8.57859775e-02, 5.52252270e-02, 8.55132714e-02, -9.58903059e-02, + 1.34693608e-01, -1.06799692e-01, 8.42991620e-02, -7.29722008e-02, + -3.16112302e-02, -1.33186892e-01, -3.01891286e-02, -1.23371415e-01, + 2.47585457e-02, 9.32596158e-03, -8.83968234e-01, 3.25804874e-02, + -1.36455193e-01, 5.22013418e-02, 7.06611201e-02, 2.37159967e-01, + 8.91016126e-02, -2.99384054e-02, -1.87199526e-02, 2.26440327e-03, + -2.86240019e-02, 1.45185545e-01, 1.16965972e-01, -3.56947407e-02, + 1.23156169e-02, -7.89485201e-02, 2.42000669e-01, 5.18658943e-02, + 1.23368412e-01, -1.33664429e-01, -5.28409593e-02, -1.49152186e-02, + -7.04273768e-03, 4.40860493e-03, -9.83215496e-02, 4.80295867e-02, + 1.94351804e-02, 9.06426013e-02, -1.53706968e-02, 5.25671393e-02, + 7.91645125e-02, 2.07034983e-02, 7.01766312e-01, 1.80371366e-02, + -1.23936214e-01, 3.77195626e-02, -9.12794545e-02, -5.07131638e-03, + -5.29812515e-01, 2.14143340e-02, 7.39834160e-02, -7.41400057e-03, + 6.01231493e-02, 7.33223036e-02, 1.53090119e-01, 1.70849040e-01, + -4.07319590e-02, -4.82572988e-03, 2.90474202e-02, 4.86663431e-02, + -1.04625016e-01, -8.78558010e-02, -4.20586728e-02, 3.25415730e-02, + -7.98427388e-02, -1.87286697e-02, -2.41589937e-02, 8.27142298e-02, + 4.98151742e-02, 5.18750064e-02, 3.29844393e-02, -3.68855074e-02, + 8.92899334e-02, -4.56710160e-03, 2.97619849e-02, -7.04986304e-02, + -1.08520284e-01, -1.11746334e-01, -7.71660432e-02, 8.49972386e-03, + 2.37586319e-01, -1.24874555e-01, -1.93465855e-02, 8.61749053e-02, + 8.92835855e-02, -8.50701034e-02, -3.94290835e-02, -2.06257198e-02, + -1.92599650e-02, 2.76334919e-02, 6.91349506e-02, -2.97043975e-02, + 8.56228359e-03, 5.08836359e-02, -1.17148422e-02, -2.16257907e-02], dtype=float32), + kernel: Array([[ 0.16482222, -0.22908561, 0.5409443 , ..., 0.52288747, + 0.45133135, -0.3148359 ], + [ 0.04791335, -0.19366321, 0.02344365, ..., 0.25754645, + 0.15889917, -0.26091245], + [ 0.21047738, 0.00189075, -0.053845 , ..., -0.10111952, + -0.42335704, -0.05204519], + ..., + [-0.07512332, 0.07116441, -0.20204972, ..., -0.3296345 , + 0.2802311 , -0.00193282], + [-0.01015196, 0.04124757, 0.08690923, ..., 0.38333163, + -0.18434097, -0.01834725], + [ 0.0501605 , 0.7881087 , 0.21703115, ..., -1.006583 , + -0.5440753 , 0.06839034]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([ 7.00707082e-03, -2.13932544e-01, -4.38091122e-02, -1.80086847e-02, + 1.56905562e-01, 3.25518586e-02, -6.51721493e-04, -3.70227383e-04, + -9.39783454e-03, -1.21371552e-01, 6.61288295e-03, -3.71838771e-02, + 1.69692468e-02, -1.13847349e-02, -1.44199710e-02, -1.44837052e-03, + 1.02954609e-02, 7.14356266e-03, -2.28633024e-02, -9.65952054e-02, + -7.08255393e-05, 4.81518777e-03, -1.34146630e-05, 3.52103100e-03, + -1.46182910e-01, -2.66583692e-02, -2.81199813e-02, 2.20761336e-02, + 8.75673592e-02, 1.68776400e-02, 7.70330988e-03, -7.63349526e-04, + 1.69382174e-03, -5.38112130e-03, -2.32931431e-02, 1.33482378e-03, + 5.63706718e-02, -1.53671801e-01, -5.60674816e-04, 1.17010390e-02, + 4.67808433e-02, -1.00002938e-03, 7.54884481e-02, 9.78838280e-03, + -2.51426380e-02, -1.93388745e-01, -1.30395563e-02, -3.22579712e-01, + -1.23171741e-02, -2.32061837e-02, -5.90802915e-02, 6.99069537e-03, + 1.86737273e-02, -1.55908689e-01, 5.42053096e-02, -1.68696530e-02, + -3.76862218e-03, 1.75712314e-02, -2.17723008e-02, -1.39841557e-01, + -1.11271013e-02, -8.20189528e-03, -1.91661697e-02, -2.04071011e-02, + 5.92404092e-03, 2.63817403e-02, 2.42956262e-02, 6.90133497e-02, + -7.82949443e-04, 6.91150129e-03, -6.76295608e-02, -8.51555727e-03, + -1.79697247e-03, 4.60966006e-02, -4.87765044e-01, 3.24066654e-02, + -7.08101019e-02, -2.59833187e-02, -4.40584077e-03, 3.93147208e-03, + 2.80292213e-01, -3.62146534e-02, 1.14013357e-02, 1.44051407e-02, + 8.79329368e-02, -3.40585620e-03, 1.61548674e-01, -4.32909727e-02, + 5.22663780e-02, -8.88153631e-03, -5.72965667e-03, 1.60463457e-03, + -1.44764967e-03, -4.41680253e-02, -6.15086630e-02, 2.85431091e-03, + -5.11140227e-02, 3.58232483e-02, 3.14552360e-03, -8.34190473e-02, + -6.30146563e-02, 3.52296643e-02, 1.75646115e-02, -2.81428844e-01, + 1.00661494e-01, -1.63044967e-02, -1.77673250e-03, -2.50048609e-03, + 9.31597501e-02, -5.21637946e-02, -1.21001536e-02, -3.70397139e-03, + -1.89869609e-02, 4.85324338e-02, -1.78271660e-03, -1.17816813e-02, + 6.07670583e-02, -6.83511868e-02, -7.59183466e-02, 2.76686642e-02, + -2.92036515e-02, 7.23615754e-03, -1.44946417e-02, 6.89900964e-02, + -1.19854533e-03, -3.26084495e-02, -3.72714689e-03, 8.10674950e-02], dtype=float32), + scale: Array([ 0.01037673, 0.20990598, 0.17576548, 0.21721113, 0.29847282, + 0.69343257, 0.00376189, -0.01851019, 0.03572948, 0.19570692, + -0.02536977, 0.42987674, 0.01018215, 0.15321532, 0.48395535, + -0.00318469, 0.16068086, 0.01112385, 0.36739224, 0.4355385 , + -0.01194211, -0.00359833, -0.086844 , 0.00597237, 0.5273556 , + 0.03060911, 0.6654633 , -0.05266754, 0.36045882, 0.10801558, + 0.29989985, -0.04426623, 0.00284036, 0.02850281, 0.08200129, + 0.00828265, 0.4786764 , 0.32044822, 0.6007385 , 0.39719838, + 0.38223657, -0.02341203, 0.02303379, 0.25033328, -0.08627673, + 0.5241369 , 0.544966 , 0.3724897 , -0.32547915, -0.24456456, + 0.32741693, 0.11249542, 0.07124001, 0.5731743 , 0.22054233, + 0.68054676, -0.04105046, 0.07573894, 0.05682403, 0.4511571 , + 0.51980895, 0.05435016, 0.06856246, 0.16396722, 0.00960637, + 0.06122901, 0.23258007, 0.3978171 , 0.23067153, -0.0107969 , + 0.517839 , 0.03140327, 0.02615155, 0.70969087, 0.1419986 , + 0.37057772, 0.54277253, 0.18082201, 0.0020515 , -0.11159092, + 0.43598834, 0.37816525, 0.586952 , 0.35433614, 0.50891733, + 0.00721514, 0.27681717, 0.38934416, 0.4676223 , 0.18506213, + 0.18672895, 0.01464603, -0.00548072, 0.72455907, 0.26698393, + -0.00541573, 0.8330279 , 0.64272374, 0.00286223, 0.52491736, + 0.5533957 , 0.3439355 , 0.32210883, 0.56023586, 0.5354366 , + 0.3412869 , 0.05528033, -0.15407614, 0.28940812, 0.39378768, + 0.5161475 , -0.03409952, 0.05442618, 0.390517 , 0.00593842, + 0.06129869, 0.42658743, 0.4008259 , 0.46517026, 0.30656123, + 0.2991743 , 0.03482211, 0.5670987 , 0.35471496, 0.08178211, + 0.0831662 , 0.00438205, 0.29678687], dtype=float32), + }, + layernorm_before: { + bias: Array([-1.00732700e-03, -3.65120143e-01, -3.34881805e-02, 4.26957064e-04, + -7.94777123e-04, -6.86906278e-03, -1.02201905e-02, -4.80189286e-02, + 1.55887365e-01, -1.59472153e-01, 5.88005059e-04, -4.74225767e-02, + 1.99526129e-03, -6.20241184e-03, 2.79833134e-02, 7.60109648e-02, + -4.03508246e-02, -5.76728676e-03, -8.25315190e-04, -5.45392111e-02, + -9.78699885e-03, -2.21987873e-01, 8.89718998e-03, 1.56212691e-02, + 2.51562544e-03, -2.31638588e-02, 2.57532690e-02, -3.37633025e-03, + -4.75819921e-03, 2.06635101e-03, 4.15294990e-02, -1.22549701e-02, + 9.92366020e-03, 1.33297034e-02, -6.88615516e-02, -2.13969056e-03, + -9.33400914e-03, -1.47152513e-01, 1.47844050e-02, 2.68603060e-02, + 1.06496513e-01, -1.40187237e-03, 2.46544197e-01, -1.18307155e-02, + 1.22905318e-02, 7.02085486e-03, 3.16104256e-02, 1.29607348e-02, + -1.19067011e-02, 5.89615945e-03, 1.45399803e-02, 7.92596638e-02, + 8.44230223e-03, -5.16842827e-02, 2.47099735e-02, 7.68707599e-03, + -7.68995145e-03, 3.25880721e-02, -3.79534140e-02, -1.32843507e-02, + 1.00192919e-01, 8.90694335e-02, 1.43536124e-02, -1.85336508e-02, + 4.20249486e-03, 4.04581055e-02, 1.56428572e-02, 5.54322731e-03, + 7.59444525e-03, -9.50715330e-04, -5.35057345e-03, 2.70144927e-04, + 1.31974250e-01, 2.08387990e-02, -5.05078256e-01, -4.19664755e-02, + 2.06986954e-03, 2.65133400e-02, 1.08654350e-02, -2.48142332e-02, + -2.32796840e-04, -3.82554834e-03, -3.06907180e-03, -3.17831822e-02, + -2.74789296e-02, 1.76229049e-04, -5.93862869e-03, -2.54624914e-02, + 5.30235134e-02, 3.77030671e-02, -3.30569521e-02, 6.81126863e-02, + -3.34922522e-02, -5.41602597e-02, -4.98229451e-03, -1.29764492e-03, + 1.27451606e-02, 5.30401291e-03, 8.72196928e-02, 5.26771834e-03, + -5.91456294e-02, 2.01584827e-02, -1.72205925e-01, -2.79393774e-02, + -3.19484659e-02, -3.31652700e-03, 3.22748870e-02, 1.10277152e-02, + 7.98177347e-02, -2.52419740e-01, 6.41727597e-02, 2.78249923e-02, + -1.18425563e-01, 5.46144024e-02, -1.38351889e-02, 1.97037379e-03, + 3.38066697e-01, -1.05525143e-02, -5.22397608e-02, -8.39580037e-03, + 1.55923087e-02, 2.35547218e-02, -4.73677404e-02, 7.04186484e-02, + -5.95991965e-03, 5.85895148e-04, 4.35512178e-02, -9.74338278e-02], dtype=float32), + scale: Array([-0.00892128, 0.22834587, 0.07191961, -0.03907548, -0.00996873, + 0.35814652, 0.27738327, 0.11775097, 0.07957964, 0.2968416 , + 0.01412186, 0.27139878, -0.00244969, -0.01754642, 0.31717873, + 0.24255775, 0.34514305, -0.0094078 , -0.01727619, 0.2070495 , + 0.23165925, 0.48358163, 0.08980764, 0.16053861, 0.07536255, + 0.02379382, 0.3457654 , 0.22846563, -0.1335053 , 0.01941462, + 0.26780826, 0.09502309, 0.14652625, 0.05376663, 0.22817616, + 0.20199119, 0.2713218 , 0.26872477, -0.0762281 , 0.22205614, + 0.31939635, 0.24618064, 0.17259929, 0.21590437, 0.02615085, + 0.17664434, 0.36018652, 0.12768349, 0.20409435, 0.14295201, + 0.18369466, 0.3224867 , 0.01613115, 0.11641789, 0.19955999, + 0.08757862, -0.01346728, 0.198473 , 0.29986396, 0.36929435, + 0.5931218 , 0.19321899, 0.11543731, 0.16761477, -0.00143133, + 0.07413565, 0.2747692 , 0.02911899, 0.12630646, 0.02624349, + 0.3049502 , 0.17227921, 0.35921425, 0.295701 , 0.13059482, + 0.19860806, 0.01733009, 0.24637464, 0.00986639, 0.19075045, + -0.00124613, 0.27278668, 0.01111688, 0.16424142, 0.18801676, + 0.01323086, 0.00368139, 0.03970122, 0.26844364, 0.10092615, + 0.09264332, 0.16377948, 0.22653754, 0.35064098, -0.05561429, + -0.00500794, 0.07151927, 0.0050664 , 0.2468497 , 0.10480628, + 0.17903912, 0.08759426, 0.41523352, 0.24769782, 0.11629754, + 0.03568345, 0.20900775, 0.05634574, 0.4340352 , 0.3333243 , + 0.41587278, 0.20580357, 0.3584831 , 0.2544934 , 0.01131326, + 0.34364027, 0.22543208, 0.13177043, 0.19698657, 0.21796568, + -0.03295131, 0.298913 , 0.47042173, 0.26397368, 0.02212258, + 0.01400256, 0.3371711 , 0.22812258], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-5.27199693e-02, -1.22556537e-02, -4.80708107e-02, -2.20634907e-01, + -1.40253054e-02, -4.87366840e-02, -8.70805699e-03, -1.54133514e-02, + -2.52253171e-02, -3.50418165e-02, -2.73905657e-02, -3.95800360e-02, + -3.71573754e-02, -2.81724222e-02, -1.33444788e-02, -1.28413305e-01, + -4.53637205e-02, -3.11292056e-02, -3.32505777e-02, -3.26416194e-02, + -3.45684476e-02, -1.40969567e-02, -2.91586872e-02, -2.76421569e-02, + -3.73606943e-02, -3.91720347e-02, -3.77491489e-02, -1.35208545e-02, + -2.09600553e-02, -4.55045812e-02, -5.19248992e-02, -2.19188482e-02, + -1.82508994e-02, -6.17397487e-01, -3.09374705e-02, -5.18258326e-02, + -2.61735450e-02, -3.89497764e-02, -1.46954479e-02, -1.18027776e-02, + -2.21383367e-02, -3.07989083e-02, -6.69780612e-01, -1.84117928e-02, + -4.12670113e-02, -8.98960326e-03, -3.24395895e-02, -1.61407497e-02, + -4.12646681e-02, -5.37053635e-03, -3.38159055e-02, -3.29089276e-02, + -1.68235525e-02, -3.26590054e-02, -2.52911523e-02, -4.61294129e-02, + -5.67609370e-01, -2.84224618e-02, -4.13801000e-02, -3.34288217e-02, + -3.53450403e-02, -3.29895802e-02, -4.02423888e-02, -2.35914946e-01, + -4.60832685e-01, -3.70765813e-02, -2.73666903e-02, -2.54514460e-02, + -1.44379437e-02, -4.47799824e-02, -4.39045168e-02, -4.60769348e-02, + -1.07184583e-02, -2.94511896e-02, -3.48258428e-02, -4.94903252e-02, + -4.71179605e-01, -3.04534864e-02, -2.14969758e-02, -2.71185916e-02, + -2.13122219e-02, -2.35180482e-02, -3.83645929e-02, -2.21622270e-02, + -2.68899292e-01, -3.47404443e-02, -3.07555981e-02, -2.66567357e-02, + -2.49998886e-02, -4.09555882e-01, -3.07542346e-02, -1.73584837e-02, + -3.32305208e-02, -3.31670046e-02, -3.28688100e-02, -2.02093106e-02, + -3.04487552e-02, -4.77275476e-02, -4.22687948e-01, -4.02031802e-02, + -2.18092110e-02, -3.58968377e-02, -2.14504618e-02, -3.14717107e-02, + -6.16603434e-01, -3.43027376e-02, -4.16996516e-02, -2.88160015e-02, + -3.17607224e-02, -4.91070986e-01, -4.11182418e-02, -3.82331535e-02, + -1.99981369e-02, -3.31944197e-01, -1.49236154e-02, -4.61634010e-01, + -4.73840237e-02, -1.24957161e-02, -3.81090641e-01, -1.33936629e-01, + -4.61227804e-01, -3.94086614e-02, -5.02103120e-02, -1.19854463e-02, + -2.26863697e-02, -3.57204154e-02, -2.89868433e-02, -3.16134878e-02, + -2.95881126e-02, -3.33401635e-02, -2.84984969e-02, -3.39026898e-02, + -4.84661758e-01, -3.00496016e-02, -2.10581459e-02, -3.03259753e-02, + -2.73086485e-02, -2.92658973e-02, -3.28680202e-02, -4.29259539e-02, + -1.15742823e-02, -3.06628253e-02, -3.49430703e-02, -4.13970686e-02, + -4.64555889e-01, -2.44929623e-02, -2.32402366e-02, -4.51661311e-02, + -3.42191793e-02, -1.04636485e-02, -2.76410598e-02, -2.72743534e-02, + -3.71849462e-02, -2.61768941e-02, -2.64907163e-02, -4.13466580e-02, + -2.31192280e-02, -2.66561229e-02, -2.44059507e-02, -2.19917297e-02, + -3.10962517e-02, -3.92389782e-02, -2.89020780e-02, -4.72296998e-02, + -1.47718657e-02, -4.71781231e-02, -3.74658257e-02, -3.74986231e-02, + -3.42181958e-02, -3.34076025e-02, -1.90843388e-01, -3.28308083e-02, + -1.35448873e-02, -3.84903140e-02, -3.14540192e-02, 9.85380495e-04, + -2.98315007e-02, -4.16581221e-02, -2.44194381e-02, -2.18599793e-02, + -4.54461351e-02, -4.33417000e-02, -2.37555113e-02, -2.65070703e-02, + -3.74873281e-02, -2.00963188e-02, -5.27748942e-01, -1.71016864e-02, + -2.98697986e-02, -3.35949808e-02, -4.21402045e-02, -4.34817523e-02, + -3.61697711e-02, -3.83502319e-02, -2.63728369e-02, -9.91544034e-03, + -4.13909107e-02, -3.27453129e-02, -1.61904916e-02, -2.92313341e-02, + -1.82820782e-02, -3.27837802e-02, -4.96267676e-02, -2.13364977e-02, + -3.07123903e-02, -1.29112750e-02, -3.05922143e-02, -1.34589160e-02, + -1.97699457e-01, -1.26856491e-02, -3.33952643e-02, 3.23457294e-03, + -2.03445368e-02, -2.09414233e-02, -1.26234703e-02, -3.02408133e-02, + -2.34196670e-02, -1.23696076e-02, -2.40369998e-02, -2.81809513e-02, + -2.13961061e-02, -5.47036016e-03, -4.53693084e-02, -3.31229195e-02, + -1.50107536e-02, -2.91419178e-02, -2.07277276e-02, -5.49518093e-02, + -2.95787137e-02, -2.38886494e-02, -2.65640840e-02, -2.46801712e-02, + -3.84779572e-02, -1.79381035e-02, -2.02893633e-02, -3.30808014e-02, + -9.46734101e-03, -3.95300575e-02, -3.69772725e-02, -3.11741419e-02, + -4.37073298e-02, -2.90284753e-02, -3.63519609e-01, -1.85389761e-02, + -1.82310361e-02, -1.51738618e-02, -1.89224500e-02, -2.03151274e-02, + -3.81886885e-02, -2.58216038e-02, -3.81475925e-01, -2.79287975e-02, + -1.52051644e-02, -2.38528401e-02, -4.04975682e-01, -3.95358130e-02, + -4.14295375e-01, -4.38631624e-02, -4.29423749e-02, -3.17317769e-02, + -3.27443145e-02, -2.63639372e-02, -3.31052691e-02, 4.43929108e-03, + -3.27373035e-02, -2.39749160e-02, -3.15703414e-02, -5.37155151e-01, + -8.51269998e-03, -3.57643776e-02, -3.98143642e-02, -3.45447585e-02, + -2.49514170e-02, -2.87994146e-01, -5.97163796e-01, -3.36527936e-02, + -3.88488770e-02, -3.22511718e-02, -4.89461757e-02, -3.97258699e-02, + -2.26812717e-02, -7.74491072e-01, -5.23366332e-01, -4.26757559e-02, + -3.48140113e-02, -2.94348393e-02, -3.54919247e-02, -2.96318009e-02, + -1.08983302e-02, -3.76659594e-02, -1.47500113e-02, -2.93985605e-02, + -2.53436007e-02, 5.79031697e-03, -2.69875340e-02, -2.84397658e-02, + -3.33693698e-02, -4.03800532e-02, -2.46412121e-02, -2.71790083e-02, + -6.01157732e-02, -3.36032584e-02, -3.37207466e-02, -1.18618924e-02, + -4.47471261e-01, -4.70100343e-02, -1.94270499e-02, -1.41834235e-02, + -4.12988141e-02, -4.31584977e-02, -2.58529224e-02, -3.11732627e-02, + -3.12253926e-02, -4.00089920e-02, -2.82256324e-02, -1.63722355e-02, + -2.23478712e-02, -3.37487720e-02, -1.17873568e-02, -3.71523947e-02, + -3.02173067e-02, -2.47102734e-02, -1.75078567e-02, -8.67900550e-01, + -4.19426337e-02, -2.89364811e-02, -3.62409726e-02, -1.40926614e-02, + -3.56511444e-01, -3.45515162e-02, -5.06763607e-02, -2.30538756e-01, + -3.77509855e-02, -3.81864123e-02, -3.19356658e-02, -2.84604561e-02, + -4.72619906e-02, -1.65800638e-02, -4.46446165e-02, -1.26604997e-02, + -3.79219428e-02, -3.04708660e-01, -3.42885070e-02, -4.41885963e-02, + -3.81699838e-02, -3.03119402e-02, -2.27376912e-02, -1.89236663e-02, + -4.71340455e-02, -1.56659111e-02, -3.82196195e-02, -3.61527875e-02, + -2.05292553e-02, -3.56830619e-02, 3.07818269e-03, -7.64196873e-01, + -2.73544770e-02, -5.15536249e-01, -1.76008549e-02, -6.18065059e-01, + -3.16055007e-02, -4.54795063e-02, -2.52228677e-02, -1.78863723e-02, + -3.71731035e-02, -3.15613933e-02, -3.33131254e-02, -1.16747636e-02, + -1.52025633e-02, -3.18961823e-03, -2.85964645e-02, -3.42308469e-02, + -6.01763930e-03, -1.44698163e-02, -3.33204381e-02, -1.46625296e-03, + -1.37048895e-02, -2.66424436e-02, -4.91908789e-02, -3.08587961e-02, + -2.29315758e-02, -3.17254931e-01, -1.58655606e-02, -8.84249713e-03, + -1.90528110e-02, -4.16713282e-02, -3.05165984e-02, -3.23779024e-02, + 1.71483532e-02, -1.71599649e-02, -1.27734728e-02, -3.09068114e-02, + -3.11801564e-02, -3.09810303e-02, -3.29174399e-02, -4.35743593e-02, + -3.38067338e-02, -2.26877723e-02, -3.14960629e-02, -2.18670722e-02, + -2.10534912e-02, -2.37145834e-02, -3.18468064e-02, -3.17289941e-02, + -3.52053195e-02, -3.62415873e-02, -3.78505550e-02, -8.54883492e-01, + -6.25598431e-01, -2.91614741e-01, -4.12686355e-02, -3.40597779e-02, + -4.49921228e-02, -1.83722675e-01, -6.20228611e-03, -9.09710955e-03, + -4.55179624e-02, -4.24469262e-02, -4.34849598e-02, -1.18993660e-02, + -2.96284892e-02, -1.82864573e-02, -1.58976540e-02, -4.17148732e-02, + -1.04806479e-02, -2.51498409e-02, -3.05920038e-02, -2.39100829e-02, + -2.05667838e-02, -4.11339812e-02, -2.58969050e-02, -2.54813936e-02, + -3.46504189e-02, -3.05222347e-02, -3.66321839e-02, -1.69235300e-02, + -3.10963802e-02, -2.00775322e-02, -3.57228443e-02, -4.01586331e-02, + -2.64529884e-02, -4.16483805e-02, -4.06530350e-02, -2.00617034e-02, + -4.20974828e-02, -3.35381776e-02, -4.97709587e-02, -1.51588218e-02, + -4.78560105e-02, -1.61089674e-02, -2.51727998e-02, -4.55499999e-02, + -3.32853608e-02, -9.99737252e-03, -9.59883165e-03, -3.51412706e-02, + -3.77242006e-02, -1.95016433e-02, -8.43567073e-01, -1.19444542e-02, + -1.49743442e-04, -3.27145346e-02, -3.28630805e-02, -1.98399983e-02, + -1.64248962e-02, -3.85820344e-02, -1.64614365e-01, -1.02842608e-02, + -1.50990197e-02, -3.22756879e-02, -1.80973206e-02, -2.27656811e-02, + -4.72140424e-02, -3.06061078e-02, -5.52861243e-02, -3.98167334e-02, + -3.09152249e-02, -3.28242555e-02, -3.62482332e-02, -2.84483526e-02, + -7.59550452e-01, -4.58953902e-02, -2.13591140e-02, -2.78457664e-02, + -1.19789401e-02, -2.98056193e-02, -3.15071791e-02, 5.63531043e-03, + -2.39895564e-02, -3.00265819e-01, -4.11561355e-02, -2.46453267e-02, + -4.19642702e-02, -5.44620812e-01, -3.15308571e-02, -3.95948440e-02, + -2.19081920e-02, -2.35356446e-02, -1.36993527e-02, -1.34584429e-02, + -5.07402495e-02, -3.16087119e-02, -2.87754349e-02, -5.01885235e-01, + -4.85699102e-02, -4.50458750e-02, -2.57593673e-02, -2.32719798e-02, + -3.24436370e-03, -2.84067634e-02, -2.62353010e-02, -3.73448022e-02], dtype=float32), + kernel: Array([[ 0.03366351, -0.01376447, 0.00334523, ..., -0.01204627, + 0.0351628 , 0.0058507 ], + [ 0.01479752, -0.00059022, 0.02888848, ..., 0.01372632, + 0.05084606, -0.00852082], + [-0.08748913, 0.06752201, -0.02948284, ..., -0.04957195, + -0.01602156, -0.0564839 ], + ..., + [ 0.03158367, -0.0024418 , -0.01248567, ..., 0.00120247, + -0.03067757, 0.05438066], + [ 0.00710372, -0.07475934, -0.03510208, ..., -0.00362264, + -0.01655151, 0.0695975 ], + [ 0.03609428, -0.04006155, -0.0374678 , ..., -0.02961577, + -0.03186451, 0.05374904]], dtype=float32), + }, + output: { + bias: Array([-1.57430485e-01, 2.07714707e-01, -1.93849746e-02, -2.16451492e-02, + -2.94376239e-02, -1.40355468e-01, 1.01594821e-01, 9.77267139e-03, + 3.38900611e-02, 3.08969438e-01, -4.50306423e-02, 4.16773744e-02, + -5.64103350e-02, 4.65862826e-02, -1.98888164e-02, 5.25933765e-02, + 5.03250724e-03, -9.93339792e-02, -9.45335254e-02, 1.34056509e-01, + -2.08574124e-02, -2.76607983e-02, -2.45273374e-02, -2.30302606e-02, + 4.72312383e-02, 6.59246966e-02, -6.18701801e-02, -2.54601543e-03, + 8.12924057e-02, 4.04679216e-02, 7.19875395e-02, -8.80277827e-02, + 1.21765286e-01, -1.06022269e-01, 8.69699791e-02, -7.91017562e-02, + -7.66793415e-02, -9.08406228e-02, -1.50119308e-02, -1.08254775e-01, + 1.43383909e-02, 1.25184851e-02, -8.86707306e-01, 3.01385261e-02, + -1.39143646e-01, 1.25063658e-01, 5.04919551e-02, 2.49843255e-01, + 5.21366931e-02, -3.97394188e-02, -1.14855161e-02, -1.06901973e-02, + -2.97776181e-02, 1.87453896e-01, 8.37227628e-02, -1.05806980e-02, + -2.71216761e-02, -8.99265856e-02, 2.36086398e-01, 1.00231051e-01, + 1.18218437e-01, -1.38503894e-01, -4.76376042e-02, -2.26873197e-02, + -9.81818512e-03, -6.29501976e-03, -1.04944162e-01, -4.35994677e-02, + 1.57965776e-02, 8.73260945e-02, 8.73774476e-03, 5.19699082e-02, + 7.45332241e-02, -8.56231526e-03, 7.00146556e-01, 9.85437538e-03, + -9.72239524e-02, 3.94646786e-02, -9.50414538e-02, -3.36449943e-04, + -4.73328263e-01, 4.33811136e-02, 1.06056370e-01, -4.22528051e-02, + 2.30107084e-02, 1.08815312e-01, 1.21224724e-01, 1.74949020e-01, + -2.51618456e-02, -2.49235518e-02, 2.24952735e-02, 4.77656983e-02, + -1.08213760e-01, -1.28404409e-01, -3.41621116e-02, 3.01731993e-02, + -7.27362260e-02, -1.10188257e-02, -2.82108001e-02, 1.34670347e-01, + 5.28050587e-02, 5.32444753e-02, 3.20014991e-02, 5.04108816e-02, + -3.19805578e-04, -1.66307837e-02, 3.41970064e-02, -6.60975128e-02, + -1.23339623e-01, -9.24760997e-02, -8.67600888e-02, 8.86741187e-03, + 2.36912012e-01, -1.58255234e-01, -1.41610643e-02, 8.07640478e-02, + 2.81964112e-02, -8.16835240e-02, -1.64333303e-02, -2.98798494e-02, + -7.30156377e-02, 2.87766512e-02, 1.21695735e-01, -5.90394363e-02, + 2.60551348e-02, 5.38477339e-02, -2.07343251e-02, -6.99748248e-02], dtype=float32), + kernel: Array([[ 0.02169053, 0.02691332, -0.00372989, ..., -0.00424074, + 0.03661725, 0.08953198], + [-0.01738822, 0.01985764, 0.03393567, ..., -0.00211225, + 0.01224655, 0.03731351], + [ 0.00689401, 0.02467112, 0.01469113, ..., -0.01584834, + -0.03345798, 0.01720266], + ..., + [ 0.01760239, 0.00760174, 0.01594303, ..., 0.00213138, + 0.00272581, -0.03686259], + [ 0.01555281, -0.08100978, 0.04178749, ..., -0.0257984 , + -0.04736535, -0.03686064], + [-0.01932741, 0.05431129, -0.04024049, ..., -0.00516084, + -0.01021444, 0.02485615]], dtype=float32), + }, + }, + }, + }, + 5: { + attention: { + attention: { + key: { + bias: Array([ 0.15380459, -0.23325492, 0.5298237 , -0.25394794, -0.1792779 , + 0.2635233 , -0.2518766 , 0.51324993, 0.44011664, -0.15511836, + 0.1065289 , -0.18103923, -0.24789087, 0.01313723, -0.00924131, + 0.2846446 , -0.02708824, 0.01480816, 0.12186062, -0.18671922, + 0.47530514, -0.44773427, -0.42629737, -0.25819948, 0.15633082, + -0.48324105, -0.2818154 , -0.61912835, -0.29742548, 0.28758916, + 0.07666907, 0.27997455, -0.5159455 , -0.00791557, 0.40676025, + -0.13460469, 0.5236954 , -0.5017536 , 0.5868691 , 0.01062866, + -0.10968728, -0.03771151, -0.33758804, -0.53288084, -0.6708519 , + 0.60964614, 0.14677066, -0.52277976, -0.00805334, -0.3969676 , + 0.5209981 , 0.34777972, 0.5903865 , -0.7917442 , -0.64988923, + -0.63839513, -0.74101704, 0.11602865, 0.8335415 , 0.2829763 , + -0.41182905, -0.3442916 , -0.47213373, 0.5067961 , 0.05412646, + 0.11681292, 0.26332203, -0.25662488, 0.4626192 , -0.2888915 , + 0.3559262 , 0.43773538, -0.4578352 , 0.32878822, -0.2354647 , + -0.39989737, -0.21824938, 0.23775657, 0.28870043, -0.35698473, + 0.2748483 , 0.05920567, -0.09002119, -0.35282713, 0.20067027, + -0.54461175, -0.23671876, -0.26704383, -0.06359763, 0.01624733, + -0.20986806, -0.0469152 , -0.21808156, 0.10429718, 0.28616387, + -0.0075318 , 0.41809165, 0.21852025, -0.68503106, 0.15126543, + -0.5385752 , -0.6387902 , -0.09934053, 1.1725092 , 0.9802906 , + -0.24173768, -0.3697774 , 0.30808112, 1.2169337 , -0.4494893 , + 0.80163956, 0.62732154, -0.03264387, 0.67577094, -0.39557546, + -0.40416336, 0.0219003 , 0.2829875 , -0.28239524, 0.14352222, + -0.8248839 , 0.02846357, 0.13515349, -0.19850218, -0.22375287, + 0.2135571 , -0.27206707, 0.76276755], dtype=float32), + kernel: Array([[ 0.26200238, 0.13167961, -0.11595184, ..., 0.0414915 , + 0.11238165, -0.30494452], + [ 0.22360708, 0.08118725, 0.02069189, ..., 0.00547806, + -0.10924851, 0.26473263], + [-0.16265303, 0.17587142, 0.12479865, ..., -0.07595218, + 0.1396125 , -0.19043295], + ..., + [-0.26119086, -0.38923773, 0.35491627, ..., -0.7368683 , + -0.49823397, -0.6397334 ], + [ 0.13432133, -0.16010363, 0.22101043, ..., -0.06990113, + 0.25519368, 0.09977636], + [-0.12494502, -0.38168752, 0.16611753, ..., 0.36234066, + -0.0669388 , 0.3607124 ]], dtype=float32), + }, + query: { + bias: Array([-0.18479703, -0.35146004, 0.7537319 , -0.17177609, 0.01369478, + 0.18720719, -0.17976424, 0.15347035, 0.08994126, 0.22886601, + 0.18279251, -0.35825828, -0.06468849, -0.12611765, -0.028998 , + 0.9971776 , -0.18668026, -0.04904906, -0.14856422, -0.24965602, + 0.05141901, -1.1571982 , 0.05819758, 0.14917302, 0.15365446, + -0.85294926, -0.08199202, -0.9097335 , 0.15537313, -0.26021054, + 0.12598538, -0.06024593, -0.11726313, -0.20979708, 0.20502815, + 0.22255987, 0.30414093, -0.11271469, 0.4672999 , 0.00509386, + -0.13042028, 0.22407965, -0.10005572, -0.14989175, -0.4365904 , + 0.2107629 , 0.18609835, -0.72566986, 0.18099388, -0.30312923, + 0.5998623 , -0.15350953, 0.36957732, -0.70576656, -0.68291605, + -0.420833 , -0.34906906, 0.06481054, 0.58362794, 0.31369215, + -0.82122093, -0.25893688, -0.24382648, 0.19694354, 0.16259049, + -0.11133604, -0.0934663 , -0.02774331, 0.21941422, -0.01251643, + 0.12317675, 0.6149345 , 0.14704953, 0.37044725, -0.28755513, + 0.1573624 , -0.44036826, 0.1894717 , 0.12188441, -0.49551368, + 0.14024805, -0.03632037, -0.11464448, -0.6433361 , 0.15439317, + -0.32633895, 0.00914915, -0.05114478, -0.1278209 , 0.33733654, + -0.07761955, -0.3304572 , -0.59981114, 0.14654088, 0.06847122, + -0.259691 , 0.25001425, 0.05990116, -0.645737 , 0.02534364, + -0.6641346 , -0.17594564, -0.01261867, 0.5100829 , 0.06452097, + -0.5573339 , 0.13375577, 0.25303102, 0.43233076, 0.15803403, + 0.254334 , 0.03359455, -0.05677273, -0.11994728, -0.2779966 , + -0.23615459, -0.04638248, -0.03050251, 0.3847538 , 0.0142238 , + -0.26559195, -0.03787645, 0.02767896, -0.08072381, -0.01213382, + 0.15982705, -0.07372479, 0.6117539 ], dtype=float32), + kernel: Array([[-0.21102534, 0.13572834, 0.08487151, ..., -0.37551054, + -0.08974288, -0.14673388], + [-0.42533755, 0.19477351, -0.05581117, ..., -0.19027376, + 0.11326433, 0.280313 ], + [ 0.20709127, -0.02206374, 0.00242551, ..., -0.21323436, + -0.1782914 , -0.49215955], + ..., + [-0.35879493, -0.13843112, 0.20357771, ..., -0.3228759 , + -0.48138356, -0.60722345], + [ 0.409117 , 0.09512895, -0.07367483, ..., -0.79642665, + -0.3589279 , 0.36061507], + [-0.1480074 , 0.24665533, -0.10206775, ..., 0.03880555, + -0.12229352, 0.42183685]], dtype=float32), + }, + value: { + bias: Array([-0.06382484, -0.10177656, -0.01783326, 0.07238453, 0.0269923 , + 0.03472366, 0.00803764, 0.09564783, -0.02505345, 0.16271482, + 0.01969427, -0.08568954, -0.02066216, -0.03939519, -0.09977195, + -0.01792192, 0.02864774, 0.03915608, 0.02681843, 0.07789411, + -0.12741898, 0.03238543, 0.01603028, -0.11089758, -0.08953978, + 0.11356429, -0.03502819, -0.0710902 , -0.07928857, -0.07034726, + -0.00949693, -0.11724807, 0.03667774, -0.0831453 , -0.08759701, + -0.11416655, 0.05479722, 0.02118164, 0.0808093 , -0.04038942, + -0.05000678, -0.06270798, -0.00748655, -0.0456011 , 0.09870113, + 0.03644619, 0.06150641, 0.15579787, 0.05361883, 0.07626446, + -0.02815951, 0.042865 , 0.06312688, 0.07170736, 0.185382 , + 0.06517194, 0.0584194 , -0.07202493, -0.09540485, -0.08599403, + -0.03540141, -0.10775873, 0.03609048, -0.06903887, -0.10958587, + -0.09116477, -0.10746475, -0.05029682, 0.07927579, 0.10405283, + 0.12020227, 0.03288461, -0.1196027 , 0.03575619, 0.0972458 , + 0.05244093, 0.08934712, 0.12486867, 0.08823898, 0.0983156 , + 0.10495333, -0.10362058, 0.09903006, 0.09088268, -0.12305731, + 0.10125972, -0.08239742, -0.00916524, -0.02156379, -0.09864786, + 0.10193784, -0.09551246, -0.05185878, -0.09967213, -0.10428255, + 0.10507633, -0.14955677, -0.12777895, -0.07334802, 0.16633663, + -0.04327126, 0.11767196, 0.08965011, 0.17071536, -0.0584787 , + -0.09368757, -0.16054668, -0.16359754, 0.0472418 , 0.1223535 , + -0.07023636, 0.24043013, 0.15963615, -0.01664593, 0.15456486, + 0.11985014, 0.25587445, 0.17214614, -0.15644138, 0.02051393, + -0.15587153, 0.03782795, 0.0117235 , 0.04015777, 0.17542166, + -0.1388207 , 0.09689201, -0.11741312], dtype=float32), + kernel: Array([[-0.06625638, -0.10676572, -0.11858384, ..., 0.12592293, + 0.19599795, -0.10273615], + [ 0.19458613, -0.08272579, -0.01620658, ..., 0.07205058, + 0.00837682, 0.03907323], + [ 0.17456186, -0.00289036, -0.24084125, ..., -0.03393131, + -0.21872245, 0.21209157], + ..., + [-0.10240208, -0.01977782, -0.14529926, ..., 0.48570195, + 0.29290506, -0.28927922], + [ 0.5828709 , 0.14644368, 0.23230499, ..., -0.51522744, + -0.00893875, 0.11504804], + [-0.18426676, -0.12852795, -0.02424787, ..., -0.17730054, + -0.39859936, 0.3791855 ]], dtype=float32), + }, + }, + output: { + dense: { + bias: Array([-1.47393912e-01, 1.51925683e-01, -1.37644792e-02, -5.88501468e-02, + -2.74660420e-02, 2.99488921e-02, 6.36566132e-02, -3.54291759e-02, + -7.69458618e-03, 3.24805588e-01, -5.66644669e-02, 2.47724857e-02, + -1.12583041e-02, -9.56679955e-02, 1.36414379e-01, 3.80160026e-02, + 4.16892134e-02, -1.21877357e-01, -9.89358947e-02, 8.71492922e-02, + -5.38453227e-04, 6.59572333e-02, -4.53343103e-03, -4.99130785e-02, + 4.46245968e-02, 1.36791214e-01, -3.86863551e-03, 1.96640324e-02, + 6.48279041e-02, 6.96821213e-02, -2.48288596e-03, -8.72585997e-02, + 2.33484972e-02, -5.56293987e-02, 6.42355978e-02, -3.93086448e-02, + -5.08882850e-03, -8.52516666e-02, 1.37223052e-02, -1.85753062e-01, + 6.52511492e-02, -4.16962616e-03, -6.68387711e-01, -2.01383736e-02, + -1.15884997e-01, 1.01936243e-01, 4.85421829e-02, 2.81445622e-01, + 9.36387926e-02, -4.22131829e-02, -2.58780066e-02, -9.31250453e-02, + -8.50194786e-03, 1.86161455e-02, 1.08278200e-01, -1.33628231e-02, + -2.95319203e-02, -7.03121051e-02, 1.65721059e-01, 1.23157077e-01, + 1.26239538e-01, 8.69110152e-02, -6.91822842e-02, -6.16661599e-03, + -4.39871997e-02, 6.22018203e-02, -9.05643180e-02, 7.91532993e-02, + -6.81676865e-02, -5.36848418e-02, 2.38579530e-02, 2.05037221e-02, + 1.06868640e-01, -3.20243724e-02, 4.63792652e-01, 2.49073096e-02, + -1.17389895e-01, 3.09424363e-02, -7.55871981e-02, -5.58793545e-02, + -3.00298661e-01, 3.98461558e-02, -5.11545166e-02, -3.12989019e-02, + 6.69162199e-02, 7.48850405e-03, 8.32765177e-02, 1.84600934e-01, + 3.32334414e-02, 4.34078313e-02, 3.89995314e-02, 1.41597819e-02, + -9.58840549e-02, 1.53720230e-02, -5.27216345e-02, 6.26715794e-02, + -1.02486372e-01, -5.28356284e-02, -1.23486230e-02, 9.72601771e-02, + -4.58076522e-02, 1.89271364e-02, 6.57725483e-02, 7.59906247e-02, + -7.66134709e-02, -4.67620529e-02, 6.40749931e-02, 2.11690012e-02, + -8.14582258e-02, -7.82966167e-02, 3.28965578e-03, -1.84342973e-02, + 1.21189028e-01, -2.12921515e-01, -4.28700894e-02, -5.90704009e-02, + 1.14371911e-01, -1.21928066e-01, -3.76502164e-02, -3.38824801e-02, + -6.17372356e-02, -1.77653898e-02, 1.73333704e-01, 8.95729363e-02, + -6.20543491e-03, 2.95204688e-02, 2.87007019e-02, -4.14825082e-02], dtype=float32), + kernel: Array([[ 0.02049902, -0.1422533 , -0.4065625 , ..., 0.36780554, + 0.4333429 , 0.19315735], + [-0.03959692, 0.25549796, -0.06924722, ..., -0.55895936, + -0.27878842, 0.00849131], + [ 0.14230408, 0.490068 , -0.18883815, ..., -0.30614755, + -0.30285606, 0.07531235], + ..., + [ 0.0692085 , 0.07044961, 0.61252284, ..., -0.23293641, + -0.15725587, -0.13308391], + [ 0.19255584, -0.04722506, 0.23870187, ..., 0.26104337, + -0.1339734 , -0.28957063], + [-0.11599459, 0.07947025, -0.35309428, ..., -0.22296615, + 0.19523388, 0.4148677 ]], dtype=float32), + }, + }, + }, + layernorm_after: { + bias: Array([-1.27752021e-01, -1.48293898e-01, -1.61956966e-01, 2.83768959e-03, + -1.43960312e-01, 4.62501980e-02, 1.55017406e-01, -1.05546750e-02, + 1.24534264e-01, 3.37903023e-01, -4.59039863e-03, -1.85061380e-01, + 1.73028678e-01, 2.01043487e-02, 2.81281001e-03, 9.39656422e-02, + -1.36595475e-03, -2.31539598e-03, -1.25028998e-01, 2.32474544e-04, + -2.05569100e-02, -1.70720860e-01, 9.77594405e-03, 1.18461205e-02, + 2.47885808e-01, -1.36246130e-01, -1.46172762e-01, -6.74006566e-02, + 1.35634257e-03, -6.88059255e-02, -1.38106868e-01, -3.34430300e-02, + 8.12808890e-03, 9.31725744e-03, 2.75329649e-02, -1.16437320e-02, + -1.76476855e-02, -1.26601487e-01, -3.02024838e-02, -5.67115694e-02, + 9.60008055e-02, 8.26168135e-02, -5.39703900e-03, 3.32480557e-02, + -1.42446086e-01, 2.15325981e-01, 3.59231502e-01, 6.70243725e-02, + 5.95594849e-03, 2.55793538e-02, 1.20438235e-02, -1.38192281e-01, + -9.54014882e-02, -6.51674941e-02, 1.10238351e-01, -9.56654698e-02, + 9.43305194e-02, -7.60812163e-02, -7.58638158e-02, 1.66560858e-01, + 1.86092202e-02, 1.56081989e-01, 5.11778006e-03, -4.85944413e-02, + 8.28391612e-02, 5.49042262e-02, -3.91565077e-02, -2.07471265e-03, + 3.06699127e-02, 3.01516578e-02, -3.39646004e-02, -2.84999385e-02, + 7.43485242e-02, 1.12809241e-02, -1.99580923e-01, -5.58756031e-02, + 2.26514554e-03, -5.92570044e-02, -5.33170663e-02, 4.89857085e-02, + 9.51478072e-03, -2.19953805e-02, -3.24405394e-02, -1.16561711e-01, + -7.09283957e-03, -4.68143495e-03, 5.30386604e-02, -1.35984283e-03, + 7.81681612e-02, -5.89879602e-03, -1.73093099e-03, -6.07467350e-03, + -3.00015986e-01, -3.05419937e-02, -1.54022276e-02, -6.91811442e-02, + -5.47736287e-02, 2.07600649e-02, 1.68445095e-01, -1.03444420e-02, + -1.04541652e-01, 5.39848953e-02, -2.24056616e-01, 1.65632993e-01, + -1.40168756e-01, -4.03944461e-04, -4.72264551e-02, 1.25392643e-03, + 6.72694370e-02, -2.77178794e-01, -4.07134034e-02, 8.86196047e-02, + 4.30512354e-02, -1.80928916e-01, -2.78730467e-02, 5.36768697e-02, + 2.03331217e-01, -2.78739282e-03, 1.01259798e-01, 6.90696239e-02, + -9.43698585e-02, 7.54763791e-03, 1.15448259e-01, 2.77357996e-01, + 1.10673435e-01, 1.47797644e-01, -9.77296499e-04, -1.15336455e-01], dtype=float32), + scale: Array([ 1.0143379 , 0.56084377, 0.61511797, 0.875825 , 0.7977206 , + 0.86066914, 0.4466555 , 0.5817147 , 0.8288987 , 1.2778934 , + 0.14663759, 0.92707175, 0.29109854, -0.5407728 , -0.00583301, + 0.80144596, 0.39370012, -0.01780795, 0.55510044, 0.8787668 , + 0.8118457 , 0.7212926 , 0.22898938, 0.00633279, 1.0344126 , + 0.4334817 , 1.2401807 , 0.45261577, 0.04603433, 0.5680823 , + 0.6459878 , 0.6229751 , 0.7747696 , 0.87726367, 0.4118688 , + 0.02244997, 0.7460163 , 0.53923315, 0.5421808 , 0.90724677, + 0.7915582 , 0.3890947 , 0.7099728 , 0.94361395, 0.44176197, + 0.708961 , 0.51787186, 1.1014241 , 0.3872233 , 0.6767877 , + 0.7677383 , 0.61901575, 0.73735505, 0.28867054, 0.81986 , + 0.68635523, 0.40406612, 0.7550665 , 0.3736026 , 0.42338955, + 0.51958174, 0.75038564, 0.10604075, 0.47208443, 0.53189516, + 0.43309155, 0.37580943, 0.48088777, 0.4441973 , 0.5321363 , + 0.5055118 , 0.98682684, 0.7596939 , 0.835741 , 0.10439195, + 0.57524043, 0.43323192, 0.71470016, 0.38151345, 0.8561705 , + 0.6086962 , 0.8476424 , 0.68529296, 0.81576556, 1.0155932 , + 0.2647536 , 0.66374 , 0.8123001 , 0.43478364, 0.7506543 , + 0.70008427, 0.01624053, 0.5848578 , 0.672363 , 0.5528674 , + 0.17722914, 0.7533496 , 0.50877273, 0.87119836, 0.50191104, + 0.70362985, 0.7272638 , 0.79176766, 1.0370609 , 0.6759593 , + -0.02001736, 0.8079056 , 0.76534927, 0.6989189 , 0.55686975, + 0.8860217 , 0.43745935, 0.6700682 , 0.90106094, 0.05214773, + 0.5131126 , 0.7131446 , 0.76999676, 0.6244328 , 0.87544256, + 0.54827476, 0.00163172, 0.72723454, 0.770018 , 0.62415475, + 0.49478632, 0.32370028, 0.52392954], dtype=float32), + }, + layernorm_before: { + bias: Array([ 5.46789803e-02, -1.76678717e-01, -3.22135426e-02, -1.92919653e-02, + -2.71175429e-03, -5.60209937e-02, 3.75961289e-02, 2.42887228e-03, + 2.20702612e-04, -1.47421882e-01, 1.25962566e-03, -9.49192513e-03, + 3.36418934e-02, 1.43754203e-02, 2.79071257e-02, 3.28558236e-02, + 1.25313839e-02, 5.57482056e-03, 2.42552720e-03, -4.17415686e-02, + -6.24468969e-03, -7.01998994e-02, 6.89587668e-02, -7.22786319e-03, + 1.22491186e-02, -1.73403710e-01, 4.65809219e-02, -2.40579285e-02, + 4.88756225e-04, 1.74035388e-03, 6.00601882e-02, -2.83498946e-03, + -4.87266015e-03, -6.68269256e-03, 2.96081766e-03, -4.06322069e-02, + -5.13824681e-03, -5.18148541e-02, -6.03306864e-04, 3.28535922e-02, + 5.79561479e-03, -5.60164126e-03, 2.76499480e-01, -5.52190188e-03, + 1.08755060e-01, 2.78750118e-02, -1.74138229e-03, -4.13134582e-02, + -1.12804314e-02, 1.25732794e-02, -2.00521722e-02, 3.98424082e-02, + 1.83249754e-03, -1.12617195e-01, 8.08609067e-04, 6.40430721e-03, + 4.37817909e-03, 2.31746268e-02, -2.95173544e-02, 5.21452923e-04, + 5.27834482e-02, 4.03346457e-02, -6.75705774e-03, -1.17142848e-03, + 2.27539521e-02, 8.85420591e-02, -6.18154998e-04, 2.34325184e-03, + 2.68001352e-02, 3.51031497e-03, -3.19339558e-02, 1.70021728e-02, + 5.05615212e-02, -3.64386290e-02, -5.41854799e-01, -3.14815044e-02, + -8.37918895e-04, -5.96139133e-02, -5.71792945e-03, 2.48509813e-02, + 5.07591176e-04, -9.80351865e-03, -4.27319445e-02, -4.39409167e-02, + -2.25920565e-02, 5.03404811e-03, 5.27535053e-03, -2.32479889e-02, + 8.92616529e-03, 2.46615335e-02, -1.04511445e-02, 4.72732633e-02, + 2.06987560e-03, -4.62427586e-02, 2.54220869e-02, -2.62153684e-03, + -1.93837471e-02, 3.46671464e-03, 9.97190922e-02, 5.05104400e-02, + -1.72052216e-02, 5.25105570e-04, -1.76775366e-01, 3.32803018e-02, + -2.44557224e-02, -1.54652279e-02, -1.59627572e-02, 3.77600156e-02, + -2.54226532e-02, -1.87684055e-02, -8.18012469e-03, 6.22503310e-02, + -9.38802143e-04, 4.75377999e-02, -9.10606887e-03, 2.95884646e-02, + 1.39154634e-02, -4.97860415e-03, 5.38583286e-03, 2.14410182e-02, + -7.47090429e-02, 2.04133317e-02, 1.69800371e-02, 5.70605770e-02, + -5.42162033e-03, -4.40403894e-02, 3.08437683e-02, -1.20276608e-01], dtype=float32), + scale: Array([ 3.09920963e-02, 1.43276036e-01, 1.49911180e-01, 2.15380549e-01, + 8.57652703e-05, 2.60982364e-01, 3.85197014e-01, -2.69828960e-02, + -3.57246934e-03, 1.77718669e-01, 1.78545862e-02, 2.83004701e-01, + 4.92948219e-02, 2.25244924e-01, 1.80071101e-01, 3.60851258e-01, + 2.13657439e-01, 6.50591101e-04, -1.27726868e-02, 8.39237347e-02, + 1.44251525e-01, 3.37201595e-01, 2.32133791e-01, 3.55153927e-03, + 1.30111605e-01, 3.19995940e-01, 1.58201158e-01, 1.93055466e-01, + 1.13361664e-02, 1.06982812e-02, 2.68632799e-01, 1.73866183e-01, + 2.49319479e-01, 2.33700909e-02, -3.93300736e-03, 1.54559046e-01, + 7.23777190e-02, 2.99003601e-01, 2.09477082e-01, 1.43370494e-01, + 2.27352500e-01, 2.16101669e-02, 1.23486273e-01, 3.68611127e-01, + 2.57967263e-01, 3.01627338e-01, 2.64851768e-02, 3.02257270e-01, + -8.79388452e-02, 1.15935393e-01, 3.20747286e-01, 3.84979308e-01, + 2.91317515e-02, 3.35988641e-01, 1.00166025e-03, -2.34178808e-02, + 1.65450841e-03, 9.69321802e-02, 1.97463453e-01, 2.78218836e-02, + 3.64540905e-01, 4.77716655e-01, -1.46322930e-02, 3.61578852e-01, + 2.67611772e-01, 1.56464770e-01, -1.12778105e-01, 3.55959013e-02, + 2.19024077e-01, 5.52458549e-03, 2.59881735e-01, 2.20885038e-01, + 2.75174618e-01, 2.36324608e-01, 1.31579578e-01, 2.25922078e-01, + 2.05405220e-01, 2.81007886e-01, -3.42874392e-03, 1.95928365e-01, + 1.01560242e-02, 1.41977131e-01, 2.52693713e-01, 2.46991277e-01, + 1.85481176e-01, -2.18331348e-02, -1.25332624e-02, 1.17323473e-01, + 6.74787387e-02, 1.07466944e-01, 1.33154258e-01, 1.68752119e-01, + -5.74604236e-03, 2.33412385e-01, 1.67403102e-01, -5.18426392e-03, + 1.57531217e-01, 1.38286993e-01, 3.68412942e-01, 2.31308982e-01, + 1.48976266e-01, 7.62487238e-04, 3.20091337e-01, 1.83057398e-01, + 2.31623769e-01, 2.69163251e-01, 1.83478698e-01, 3.07111442e-01, + 3.13205153e-01, 2.52505481e-01, 3.30679387e-01, 3.39435637e-01, + 9.45753418e-03, 4.02425140e-01, 6.86238613e-03, 2.50461936e-01, + 3.17101419e-01, 4.53172848e-02, 6.39474615e-02, 2.83503503e-01, + 1.75329089e-01, 1.88262612e-01, 1.82763785e-01, 3.37596089e-01, + 3.82400723e-03, 2.11784661e-01, 2.38452464e-01, 3.09545726e-01], dtype=float32), + }, + moe: { + shared_experts: { + intermediate: { + bias: Array([-3.9847344e-01, -3.6854276e-01, -4.5092517e-01, -4.3007284e-01, + -4.8035741e-01, 1.8460776e-03, -4.3567291e-01, -4.2967001e-01, + -3.3667952e-01, -4.1538629e-01, -3.6099383e-01, -4.3359479e-01, + -4.4380239e-01, -5.5524331e-01, -3.6870196e-01, -2.4764657e-01, + -4.7775912e-01, -4.1283187e-01, -3.4256911e-01, -2.0356853e-01, + -4.6999124e-01, -5.6414276e-01, -3.3079070e-01, -3.8981372e-01, + -4.5676839e-01, -8.1164964e-02, -2.3958890e-01, -4.1569802e-01, + -3.3400738e-01, -3.1094922e-02, -4.0657026e-01, -3.2647237e-01, + -5.5474406e-01, -5.1273751e-01, -4.6838763e-01, -4.5339519e-01, + -3.2953310e-01, -4.2982268e-01, -4.1318712e-01, -4.8618275e-01, + -3.3651689e-01, -3.8407576e-01, -3.7777162e-01, -4.4129571e-01, + -2.5866541e-01, -4.4932291e-01, -4.2343855e-01, -3.6280292e-01, + -3.8582063e-01, -3.7247223e-01, -7.4966386e-02, -3.5159770e-01, + -4.0223166e-01, -1.5560038e-01, -3.4897429e-01, -3.5570183e-01, + -2.3108204e-01, -3.0199406e-01, -5.0925249e-01, -4.2046854e-01, + -5.7922322e-01, -4.1019771e-01, -4.2622930e-01, -4.4688463e-01, + -4.6536162e-01, -3.6277786e-01, -2.4544826e-01, -3.0414537e-01, + -4.6335497e-01, -3.0622986e-01, -3.6244193e-01, -1.3396668e-01, + -4.9729797e-01, -4.0432706e-01, -1.7137310e-01, -1.5809147e-01, + -4.3105653e-01, -2.7127942e-01, -3.0011219e-01, -2.2747953e-01, + -3.7861067e-01, -2.8296992e-01, -4.9367398e-01, -1.4615911e-01, + -3.4952867e-01, -2.4669057e-01, -4.6957275e-01, -3.5128793e-01, + -3.9153334e-01, -4.2859459e-01, -4.0044680e-01, -4.3775627e-01, + -3.7979439e-01, -5.2443427e-01, -1.5687215e-01, -4.4120738e-01, + -4.7335675e-01, -3.9626312e-01, -3.6113608e-01, -8.2696810e-02, + -2.6075220e-01, -1.9168627e-01, -4.8281315e-01, -1.7972828e-01, + -5.2597886e-01, -2.7033496e-01, -5.4604769e-01, -3.4184673e-01, + -4.3216091e-01, -4.4862813e-01, -2.6511189e-01, -4.5229965e-01, + -3.8147661e-01, -3.2180479e-01, -3.1602433e-01, -4.3689847e-01, + -5.1146489e-01, -3.5368246e-01, -3.7338638e-01, -3.0413544e-01, + -2.4527813e-01, -4.7685048e-01, -2.5364643e-01, -3.0931181e-01, + -3.2056129e-01, -3.0224264e-01, -3.6615273e-01, -3.6344633e-01, + -4.7639742e-02, -4.2077422e-01, -4.8495576e-01, -3.2529086e-01, + -3.9314100e-01, -3.9539710e-01, -4.4750559e-01, -2.1847221e-01, + -3.4810659e-01, -3.5670453e-01, -4.8384792e-01, -4.6719027e-01, + -2.9682928e-01, -4.2099711e-01, -4.7452945e-01, -4.3598908e-01, + -4.5502350e-01, -4.7916898e-01, -5.1886410e-01, -2.9352316e-01, + -4.6964645e-01, -4.7157702e-01, -2.2555082e-01, -3.9752248e-01, + -7.7490084e-02, -1.6522288e-01, -4.9982330e-01, -4.2842680e-01, + -3.2544607e-01, -3.0850667e-01, -4.0218577e-01, -4.1208234e-01, + -4.0038344e-01, -3.7057555e-01, -2.2373217e-01, -3.4376973e-01, + -3.1466734e-01, -5.2123147e-01, -5.3340238e-01, -2.8239700e-01, + -3.3951777e-01, -4.2198521e-01, -4.4805121e-01, -2.5464025e-01, + -1.5028036e-01, -4.1943482e-01, -4.2202288e-01, -3.9689219e-01, + -3.7402752e-01, -1.6487882e-01, -3.0544364e-01, -4.0235436e-01, + -4.1456652e-01, -3.5549703e-01, -4.5845604e-01, -3.7517163e-01, + -4.1666454e-01, -2.5526059e-01, -4.9071726e-01, -3.1533489e-01, + -3.4268734e-01, -4.6997485e-01, -3.2526866e-01, -4.7221246e-01, + -3.2835072e-01, -5.4744834e-01, -4.1387323e-01, -2.6110759e-01, + -4.3684092e-01, -4.7436625e-01, -3.4099048e-01, -3.7735420e-01, + -4.5077652e-01, -5.1907909e-01, -3.7325633e-01, -2.2339475e-01, + -2.8594923e-01, -4.6022233e-01, -3.7742871e-01, -3.5255402e-01, + -1.6996916e-01, -3.3114195e-01, -2.1370201e-01, -4.9245229e-01, + -4.8635155e-01, -4.3522072e-01, -3.8971159e-01, -4.3678322e-01, + -4.6781665e-01, -2.9914606e-01, -4.2043447e-01, -2.9230228e-01, + -2.6290324e-01, -2.3281787e-01, -5.2061099e-01, -3.1569782e-01, + -5.0653201e-01, -4.7192627e-01, -3.8380232e-01, -4.1764271e-01, + -3.3932805e-01, -4.6573785e-01, -4.0508986e-01, -1.6529684e-01, + -3.9876980e-01, -4.5229152e-01, -2.8774762e-01, -3.5973960e-01, + -3.5486862e-01, -2.4070273e-01, -2.9540777e-01, -2.9138458e-01, + -3.0258924e-01, -3.0566293e-01, -3.6295635e-01, -2.6400667e-01, + -3.0755830e-01, -2.9665276e-01, -3.2570139e-01, -3.8138995e-01, + -2.5602007e-01, -3.7541685e-01, -4.3872544e-01, -4.1538125e-01, + -4.2229000e-01, -3.4041786e-01, -3.0963835e-01, -4.1050071e-01, + -5.6023502e-01, -5.6029701e-01, -3.5283986e-01, -2.6442611e-01, + -3.5243306e-01, -3.4913927e-01, -2.1650518e-01, -4.0154454e-01, + -3.7275749e-01, -4.7750309e-01, -3.6822703e-01, -4.5743442e-01, + -4.8160750e-01, -4.7312003e-01, -3.9876688e-01, -5.0516546e-01, + -5.4552037e-01, -3.4587166e-01, -1.6361688e-01, -3.0174541e-01, + -3.3792850e-01, -4.7918764e-01, -4.8411125e-01, -4.3018821e-01, + -4.6167120e-01, -4.1177577e-01, -4.5708910e-01, -4.5317286e-01, + -2.9856035e-01, -3.5552701e-01, -3.3068407e-01, -4.0677980e-01, + -4.2047119e-01, -2.2582465e-01, -5.4681146e-01, -3.3787546e-01, + -4.4846484e-01, -3.3120909e-01, -3.4642521e-01, -3.9533862e-01, + -3.4406123e-01, -4.5375788e-01, -5.0441438e-01, -3.3547118e-01, + -2.8920132e-01, -1.0944715e-01, -3.5370487e-01, -3.3076125e-01, + -2.8330564e-01, -3.6274710e-01, -2.7184573e-01, -3.5993487e-01, + -4.0513000e-01, -3.0320236e-01, -2.5303522e-01, -4.7531182e-01, + -4.6395004e-01, -4.2257047e-01, -4.8910224e-01, -3.6540136e-01, + 1.3740684e-01, -3.0950975e-01, -4.7773069e-01, -3.0934384e-01, + -4.1294369e-01, -4.5992157e-01, -4.8535463e-01, -4.4464013e-01, + -3.5180721e-01, -3.6469987e-01, -4.2818370e-01, -5.0174421e-01, + -3.8951018e-01, -3.7477994e-01, -3.4967715e-01, -4.3007034e-01, + -4.4603419e-01, -3.4457707e-01, -4.6689948e-01, -3.0864301e-01, + -4.9847844e-01, -3.7400711e-01, -4.7358903e-01, -3.5766226e-01, + -2.8434107e-01, -2.7686146e-01, -3.4817812e-01, -3.6055216e-01, + -4.3933406e-01, -5.2891338e-01, -2.7604723e-01, -5.1344776e-01, + -2.9326063e-01, -5.0666833e-01, -3.3879602e-01, -4.9309245e-01, + -3.0533355e-01, -5.3628761e-01, -4.3473902e-01, -4.1327205e-01, + -1.9567779e-01, -4.3977755e-01, -4.7691789e-01, -3.4455177e-01, + -4.1720390e-01, -3.4753558e-01, -4.0889883e-01, -4.2685688e-01, + -3.7366074e-01, -1.6097990e-01, -3.0323794e-01, -4.8915049e-01, + -3.5860199e-01, -3.2890522e-01, -3.6635777e-01, -3.0436921e-01, + -4.3726847e-01, -1.7787141e-01, -3.3810359e-01, -4.7570062e-01, + -5.1493889e-01, -4.9520645e-01, -4.8514831e-01, -4.0491089e-01, + -5.2531552e-01, -4.0086076e-01, -5.0189817e-01, -4.7050786e-01, + -2.1616998e-01, -3.5684329e-01, -4.2839581e-01, -3.8793033e-01, + -2.3242010e-01, -3.5336223e-01, -3.3516747e-01, -3.7301928e-01, + -4.4247723e-01, -4.8835176e-01, -2.0324853e-01, -3.9160007e-01, + -3.3982515e-01, -2.5195253e-01, -5.4986537e-01, -3.1654030e-01, + -4.6651766e-01, -4.7360784e-01, -2.8256953e-01, -3.3791304e-01, + -3.5749507e-01, -3.0531624e-01, -4.9014235e-01, -4.1924736e-01, + -2.0831045e-02, -3.5200366e-01, -2.2583617e-01, -4.5401311e-01, + -3.8651881e-01, -4.5588291e-01, -4.9305251e-01, -1.7810422e-03, + -2.6118660e-01, -3.1825066e-01, -2.4539576e-01, -2.5099796e-01, + -5.3763103e-01, -3.6526870e-02, -2.3981811e-01, -3.0151102e-01, + -4.7348249e-01, -2.2897437e-01, -4.1141531e-01, -4.8857701e-01, + -3.2772517e-01, -5.7252842e-01, -4.5862877e-01, -3.6546475e-01, + -3.6844012e-01, -2.0294371e-01, -4.8403230e-01, -4.8141652e-01, + -1.4463230e-03, -4.9992949e-01, -3.6349079e-01, -2.1532026e-01, + -4.1290241e-01, -4.8122567e-01, -2.8661153e-01, -4.4909769e-01, + -3.6796397e-01, -2.6207677e-01, -3.9517501e-01, -4.4621187e-01, + -3.6207435e-01, -4.8050162e-01, -3.7443826e-01, -2.3039897e-01, + -3.5041386e-01, -3.9968762e-01, -3.6701423e-01, -3.8556710e-01, + -3.4414786e-01, -4.8362854e-01, -3.8469809e-01, -4.8367801e-01, + -3.7582186e-01, -4.8235339e-01, -3.0144668e-01, -3.7805808e-01, + -3.9068311e-01, -4.6765688e-01, -4.5509905e-01, -4.1643125e-01, + -4.3739045e-01, -5.5417305e-01, -3.9493406e-01, -4.5592216e-01, + -3.2413527e-01, -4.4394735e-01, -2.5142363e-01, -2.9339787e-01, + -4.6713012e-01, -2.5261605e-01, -2.6347214e-01, -3.7425596e-01, + -2.1422139e-01, -4.4427580e-01, -3.5874227e-01, -4.1653034e-01, + -3.3887890e-01, -3.0552921e-01, -2.4492532e-01, -4.0205300e-01, + -3.2229105e-01, -4.8967350e-01, -2.1930854e-04, -3.7214684e-01, + -3.0033603e-01, -4.7503376e-01, -3.7032735e-01, -3.4375238e-01, + -3.8198787e-01, -3.3686790e-01, -5.5172253e-01, -3.0320832e-01, + -1.3836212e-01, -3.9621928e-01, -3.8446391e-01, -4.3740526e-01, + -3.8315192e-01, -4.0578344e-01, -3.2919899e-01, -4.7428545e-01, + -4.4342765e-01, -4.8948532e-01, -2.3972549e-01, -3.4704226e-01], dtype=float32), + kernel: Array([[ 0.10745826, 0.48113093, -0.04176594, ..., -0.36369812, + 0.15021859, -0.05807315], + [-0.58055264, -0.5796294 , -0.19744751, ..., -0.07468567, + -0.51942426, -0.03569252], + [ 0.01138834, 0.26800624, -0.01454838, ..., -0.8072465 , + 0.5447465 , 0.36949006], + ..., + [ 0.36528957, 0.27883154, -0.3109063 , ..., -0.5552432 , + 0.23742002, 0.05708292], + [-0.09702465, 0.6521357 , 0.16876654, ..., -0.252267 , + -0.09898252, -0.03783756], + [-0.05449405, -0.4971733 , -0.02605094, ..., 0.03078811, + 0.16231588, 0.4204838 ]], dtype=float32), + }, + output: { + bias: Array([-1.16826370e-01, 1.19446300e-01, -3.04981880e-03, -1.48974463e-01, + 5.59353642e-02, 7.24940300e-02, 3.09443232e-02, -3.97321768e-02, + -1.61944836e-01, -1.69336453e-01, -5.43983765e-02, 1.14368267e-01, + -1.82648357e-02, -5.84783889e-02, 2.02381879e-01, -1.69515777e-02, + 4.80070449e-02, -1.26781985e-01, -6.17002137e-02, 3.90345929e-03, + -1.18189584e-02, 1.06815845e-01, 2.18161605e-02, -5.02868928e-02, + -1.54914372e-02, 2.92950004e-01, 3.89679447e-02, 4.18758281e-02, + 6.82518259e-02, 8.17627087e-02, 2.73418054e-02, -7.99397752e-02, + 2.45839432e-02, 1.14206569e-02, 3.90483811e-02, -3.62712778e-02, + 2.89402343e-02, -7.93813691e-02, 1.43112317e-01, -1.52107775e-01, + 1.80989935e-03, -1.87356304e-02, -5.44708908e-01, -9.66665521e-03, + -5.63648678e-02, 1.56991724e-02, -2.02740412e-02, 5.39821014e-02, + 1.57445461e-01, -2.57236585e-02, -4.68497947e-02, -2.43687704e-02, + 9.51719135e-02, 7.29279891e-02, 1.93980023e-01, 1.36569049e-02, + 6.77157715e-02, -6.52143881e-02, 2.89444447e-01, 1.13720141e-01, + 1.87333062e-01, 1.36707956e-02, -7.90089220e-02, 1.85856335e-02, + -5.92513643e-02, 7.04728812e-02, 8.90799426e-03, 1.34663433e-01, + -6.09093085e-02, -2.86442898e-02, 3.39895934e-02, -6.57519400e-02, + -6.73585013e-02, -5.29800951e-02, 5.72383285e-01, 7.09643438e-02, + -1.11088164e-01, 3.76055352e-02, -4.64858226e-02, 1.60521493e-04, + -2.10771397e-01, -1.34013280e-01, -9.27246362e-02, -2.77775893e-04, + 1.43522322e-01, 4.07825224e-02, 5.58940060e-02, 1.22946128e-01, + -1.75536349e-01, 2.35335454e-01, 8.74865577e-02, 1.05813751e-02, + 1.50552183e-01, 9.60319340e-02, 6.48361864e-03, 6.70408085e-02, + -1.03584118e-01, -6.56064600e-02, -5.26508018e-02, 1.28201723e-01, + 3.76123935e-04, -2.51332112e-02, 8.39170814e-02, 3.41377631e-02, + -2.44789012e-02, -2.70240922e-02, 5.39324842e-02, -2.40320619e-02, + -9.49936584e-02, -1.08399363e-02, 8.49590078e-02, -4.56464440e-02, + -1.49793148e-01, 6.02403209e-02, -2.90313717e-02, -4.08411808e-02, + 2.29514278e-02, -1.53553560e-01, -5.11574633e-02, 3.26815434e-02, + 1.73677616e-02, -2.31411960e-02, 1.70279935e-01, -1.79555595e-01, + -7.74298981e-02, 9.48442426e-03, 4.52735312e-02, 9.07007884e-03], dtype=float32), + kernel: Array([[ 0.49678656, -0.04703587, 0.14592949, ..., 0.14652656, + 0.4102155 , 0.43552795], + [ 0.12093473, 0.28034988, 0.05580731, ..., -1.0471447 , + 0.00430787, -0.37116638], + [ 0.4650048 , 0.25468072, 0.04566461, ..., -0.30497208, + -0.02476864, 0.01837519], + ..., + [ 0.13562478, -0.04827686, 0.8059429 , ..., -0.31421518, + 0.24673045, -0.4773905 ], + [-0.73090583, -0.39090383, 0.1884505 , ..., 0.3404823 , + 0.5766392 , 0.37859675], + [-0.57174814, 0.14778472, -0.11844227, ..., -0.2646688 , + -0.18360172, -0.04887927]], dtype=float32), + }, + }, + }, + }, + }, + }, + layernorm: { + bias: Array([ 4.99482080e-02, -4.29396838e-01, 1.81105509e-02, 4.06332692e-05, + 1.38564035e-01, -2.67413799e-02, 4.30011488e-02, -3.71168442e-02, + -5.17193019e-01, -9.74547490e-03, -6.02933764e-02, 8.66996311e-03, + 3.20597798e-01, -1.68491941e-04, 1.20097917e-04, -8.34563673e-02, + 9.34075788e-02, -6.27656728e-02, -3.47385721e-05, -4.00960334e-02, + 1.16589768e-02, 2.07135342e-02, 5.25931716e-02, -5.63471504e-02, + 1.14345705e-04, -2.30510355e-04, 2.13737637e-02, -1.45103643e-03, + 1.10674717e-01, 3.27529460e-02, 1.30938485e-01, 2.17286926e-02, + -6.83023154e-06, -9.22439562e-04, -5.39473398e-03, -4.33185026e-02, + -9.40410607e-03, -1.07025966e-01, -1.36162853e-03, 1.07634941e-03, + 2.64154427e-04, -1.61221977e-02, 1.59604579e-01, 6.28793612e-04, + 2.51234043e-03, -1.61393564e-02, -3.47336605e-02, -2.62861013e-01, + -1.80254143e-03, 1.19618746e-02, -1.74341142e-01, -2.35029664e-02, + 1.11634769e-01, -7.43765607e-02, 2.52770900e-04, 5.80298016e-03, + 1.65572765e-05, -1.99696115e-05, 8.98374594e-07, 2.72183418e-02, + 5.92560282e-05, 2.25074757e-02, -7.86222667e-02, 3.12634259e-02, + -5.19049093e-02, 7.74146914e-02, -1.94988661e-02, 3.86598520e-04, + 3.62606696e-03, 1.13942362e-02, -5.74776307e-02, 1.42280420e-04, + -7.51663325e-03, -2.08205208e-01, -3.37208897e-01, 6.53023422e-02, + -9.40454900e-02, 2.59398657e-04, -1.21128708e-01, 6.60678595e-02, + 5.90708017e-01, -3.01386812e-04, -4.50036081e-04, -4.51468586e-05, + 1.67313181e-02, 1.51245928e-04, 1.40742473e-02, 5.68821088e-06, + 1.60795447e-04, 1.45528553e-04, 6.64302963e-04, 4.23284993e-02, + -2.40530949e-02, 2.40640657e-06, -1.34843877e-02, 4.14691959e-03, + -1.12099377e-02, -5.02210781e-02, 2.03482993e-02, 2.34702706e-01, + 7.58242968e-05, -5.44176670e-04, -2.74804551e-02, 2.38151588e-02, + -4.01919484e-02, -3.65443416e-02, -7.20889375e-06, -2.77027546e-04, + -8.81933421e-02, 1.69912830e-01, 2.06964463e-02, 1.79356318e-02, + -1.09324465e-04, 6.39357895e-04, 3.84408049e-02, 6.90272776e-04, + -3.20791632e-01, 6.30142495e-06, 4.83152457e-02, -2.25914060e-04, + -3.39954719e-02, -4.70114350e-02, 2.85215719e-05, 1.06935140e-05, + -1.23181306e-02, 1.01582222e-01, 3.10630426e-02, -4.47405800e-02], dtype=float32), + scale: Array([ 6.8868792e-01, 3.1249523e-01, 9.1828078e-01, -2.3013241e-04, + 3.7263566e-01, 2.3251632e-01, 6.5610683e-01, 1.1425257e-01, + 3.4526226e-01, 1.4145033e-02, 7.2997391e-01, 8.0623746e-02, + 2.9320392e-01, -4.1230087e-05, -6.6874339e-04, 5.7173258e-01, + 6.9446439e-01, 1.2802645e+00, -6.2002582e-05, 1.1703534e-01, + 4.1233909e-01, 8.1602544e-01, 5.0138348e-01, 5.7336593e-01, + 7.1710587e-05, 1.3667290e-05, 1.9022560e-01, 6.5288413e-01, + 4.8426083e-01, 4.8938426e-01, 6.1096561e-01, 4.0365598e-01, + 2.9291890e-05, 1.8095745e-03, 4.2584249e-01, 1.0999548e+00, + 2.1008846e-01, 3.8091195e-01, 4.8326021e-03, -1.2936177e-04, + -3.4556305e-04, 6.0092252e-01, 3.3731986e-02, 2.1269962e-03, + 7.6969045e-01, 3.1354058e-01, 5.6419522e-01, 3.5723236e-01, + 1.6803181e-01, 1.7196296e-01, 6.4958000e-01, 5.5659479e-01, + 2.6263377e-01, 2.4760671e-01, -5.8708456e-04, 3.1410387e-01, + -1.7643184e-04, 2.0361553e-04, -2.0005717e-05, 6.7649913e-01, + -1.2435010e-04, 3.2834163e-01, 6.4273733e-01, 4.8385906e-01, + 8.3392429e-01, 2.8222600e-01, 1.5083265e-01, -5.6320438e-03, + 3.3680016e-01, 2.1632074e-01, 4.0965131e-01, 4.5164741e-04, + 1.3090445e-01, 7.1241659e-01, 8.9404747e-02, 4.2915529e-01, + 6.3035327e-01, -3.9689415e-04, 6.1234730e-01, 2.9691941e-01, + 1.6158453e-01, 1.7273755e-04, 1.4721676e-04, -1.1007604e-04, + 1.3981619e-01, 1.9772084e-04, 4.9367583e-01, 2.4138608e-06, + -8.6720851e-05, 3.0593758e-03, 2.4073297e-01, 1.1165603e+00, + 1.5985191e-01, 3.0023075e-04, 2.5869986e-01, 9.9650544e-01, + 2.4187727e-01, 4.2451209e-01, 4.7070310e-01, 4.9373725e-01, + -3.1352194e-04, 1.4548376e-04, 1.2102582e+00, 2.8354323e-01, + 4.2492944e-01, 3.4629729e-01, 6.0579823e-06, -1.6931001e-03, + 7.3870343e-01, 6.7030573e-01, 1.5534639e-01, 5.8719957e-01, + -1.9493036e-04, 1.6623877e-03, 7.6086062e-01, 3.6656085e-01, + 3.9590457e-01, 7.1911036e-06, 7.4346775e-01, -2.0176137e-04, + 2.7498129e-01, 6.6088724e-01, -2.2207441e-06, -7.9591937e-06, + 2.0538671e-01, 1.2334855e+00, 2.2934711e-01, 1.6446222e-01], dtype=float32), + }, + }, +}) diff --git a/password.txt b/password.txt new file mode 100644 index 0000000000000000000000000000000000000000..2e5584afeabfece9e9c34e7ef6bd8deecbe9380d --- /dev/null +++ b/password.txt @@ -0,0 +1,2 @@ +lmc-transformer: 6niAInCZxocweP6weMTKC8cobrLVHZmttuVh96OmamjSL9gjkHSdJQQJ99BFACAAAAAubb37AAASAZDO2upd +UKayDJX6eG \ No newline at end of file diff --git a/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 b/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 new file mode 100644 index 0000000000000000000000000000000000000000..1bcc83687f75b4ec26f583833434c61f21360765 Binary files /dev/null and b/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 differ diff --git a/plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed20].pdf b/plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed20].pdf new file mode 100644 index 0000000000000000000000000000000000000000..0f19af2862bfade8fc3185e64e6c98039e0cad7a Binary files /dev/null and b/plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed20].pdf differ diff --git a/plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed40].pdf b/plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed40].pdf new file mode 100644 index 0000000000000000000000000000000000000000..4507d629d9e0301419f2de186e3732d87fea879a Binary files /dev/null and b/plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed40].pdf differ diff --git a/plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed40].pdf b/plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed40].pdf new file mode 100644 index 0000000000000000000000000000000000000000..a5f88e3ddf554fbf5f83e0b5c242b6bcea822bd3 Binary files /dev/null and b/plots/enwik8/[finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice[0]-heads4-shared1-routed0-topk0-seed40].pdf differ diff --git a/plots/enwik8/[finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed20].pdf b/plots/enwik8/[finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed20].pdf new file mode 100644 index 0000000000000000000000000000000000000000..80428a409fc4226522ca32b587e6021c4403968b Binary files /dev/null and b/plots/enwik8/[finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed20].pdf differ diff --git a/plots/enwik8/[finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed40].pdf b/plots/enwik8/[finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed40].pdf new file mode 100644 index 0000000000000000000000000000000000000000..a3093a0d51c80e035d9a671aed24459b1965ef8a Binary files /dev/null and b/plots/enwik8/[finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice[0]-heads8-shared1-routed0-topk0-seed40].pdf differ diff --git a/plots/enwik8/learnable-indice0-heads16-shared1-routed0-topk0.pdf b/plots/enwik8/learnable-indice0-heads16-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..0686c3d5ac5388806e6c34b9c290f5cb4c55de54 Binary files /dev/null and b/plots/enwik8/learnable-indice0-heads16-shared1-routed0-topk0.pdf differ diff --git a/plots/enwik8/learnable-indice0-heads4-shared1-routed0-topk0.pdf b/plots/enwik8/learnable-indice0-heads4-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..58c1957a05c8d3baab5494e7be53dc9dabb0618e Binary files /dev/null and b/plots/enwik8/learnable-indice0-heads4-shared1-routed0-topk0.pdf differ diff --git a/plots/enwik8/learnable-indice0-heads8-shared1-routed0-topk0.pdf b/plots/enwik8/learnable-indice0-heads8-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..ed4f770f81771dadca776024289934cbc7a5fa4e Binary files /dev/null and b/plots/enwik8/learnable-indice0-heads8-shared1-routed0-topk0.pdf differ diff --git a/plots/imagenet/.pdf.png b/plots/imagenet/.pdf.png new file mode 100644 index 0000000000000000000000000000000000000000..d2e6ebd2abe71c3f2a6e7f987d14d16f4f050e12 Binary files /dev/null and b/plots/imagenet/.pdf.png differ diff --git a/plots/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed0+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.pdf b/plots/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed0+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.pdf new file mode 100644 index 0000000000000000000000000000000000000000..b1e1e2f69049c31d5816a3cd2753830df4c41784 Binary files /dev/null and b/plots/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed0+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.pdf differ diff --git a/plots/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.pdf b/plots/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.pdf new file mode 100644 index 0000000000000000000000000000000000000000..18a7ed57f4cef40c11f083f5b1350f350c175fd3 Binary files /dev/null and b/plots/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.pdf differ diff --git a/plots/lm1b/learnable-indice0-heads12-shared1-routed0-topk0.pdf b/plots/lm1b/learnable-indice0-heads12-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..ae94317e9e087926bd64f6fe1420dc1a0775753e Binary files /dev/null and b/plots/lm1b/learnable-indice0-heads12-shared1-routed0-topk0.pdf differ diff --git a/plots/lm1b/learnable-indice0-heads16-shared1-routed0-topk0.pdf b/plots/lm1b/learnable-indice0-heads16-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..e664f5a479bf3daa88e9f88ce53421f59065c260 Binary files /dev/null and b/plots/lm1b/learnable-indice0-heads16-shared1-routed0-topk0.pdf differ diff --git a/plots/lm1b/learnable-indice0-heads8-shared1-routed0-topk0.pdf b/plots/lm1b/learnable-indice0-heads8-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..f2cdfcadc15f77eb76ff4a0eff02649a76c29c1b Binary files /dev/null and b/plots/lm1b/learnable-indice0-heads8-shared1-routed0-topk0.pdf differ diff --git a/plots/text8/learnable-indice0-heads16-shared1-routed0-topk0.pdf b/plots/text8/learnable-indice0-heads16-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..1ed19ae9fdb7e15e46ad1facf3cd43860e3b0031 Binary files /dev/null and b/plots/text8/learnable-indice0-heads16-shared1-routed0-topk0.pdf differ diff --git a/plots/text8/learnable-indice0-heads4-shared1-routed0-topk0.pdf b/plots/text8/learnable-indice0-heads4-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..f9ffc234a3be9f69290c2d90a3d279f31763daef Binary files /dev/null and b/plots/text8/learnable-indice0-heads4-shared1-routed0-topk0.pdf differ diff --git a/plots/text8/learnable-indice0-heads8-shared1-routed0-topk0.pdf b/plots/text8/learnable-indice0-heads8-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..c4e7f027b2dc82e97c421bd032b028bd5ef3864a Binary files /dev/null and b/plots/text8/learnable-indice0-heads8-shared1-routed0-topk0.pdf differ diff --git a/plots/wt103/learnable-indice0-heads2-shared1-routed0-topk0.pdf b/plots/wt103/learnable-indice0-heads2-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..e76b27daec0f4f55abc6ac389c1b8e3a78d5bc4e Binary files /dev/null and b/plots/wt103/learnable-indice0-heads2-shared1-routed0-topk0.pdf differ diff --git a/plots/wt103/learnable-indice0-heads3-shared1-routed0-topk0.pdf b/plots/wt103/learnable-indice0-heads3-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..1d7674030b507641c37510c474af6ed24c8da3bc Binary files /dev/null and b/plots/wt103/learnable-indice0-heads3-shared1-routed0-topk0.pdf differ diff --git a/plots/wt103/learnable-indice0-heads4-shared1-routed0-topk0.pdf b/plots/wt103/learnable-indice0-heads4-shared1-routed0-topk0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..d19385d7783554b3d4919fa19d8bdb16f3771420 Binary files /dev/null and b/plots/wt103/learnable-indice0-heads4-shared1-routed0-topk0.pdf differ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..7a98417aa97d05834fa90d1c78f13742eb625101 --- /dev/null +++ b/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"] } \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..2dd5529846b022aac51a4b23ec8e3947269db0fa --- /dev/null +++ b/requirements.txt @@ -0,0 +1,38 @@ +absl-py>=0.12.0 +aqtp +chex>=0.0.7 +clu>=0.0.3 +datasets +einops>=0.3.0 +flax>=0.6.4 +ml-collections>=0.1.0 +numpy>=1.19.5 +pandas>=1.1.0 +<<<<<<< HEAD +tensorflow-cpu>=2.4.0 +tensorflow-datasets>=4.0.1 +tensorflow-probability>=0.11.1 +tensorflow-text>=2.9.0 +======= +>>>>>>> master +torch +torchvision +scikit-learn +matplotlib +tqdm +augmax +optax +dataclasses +argparse +wandb +timm +tqdm +wandb +scikit-learn +transformers +--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html +<<<<<<< HEAD +jax[cuda12_pip] +======= +jax[cuda] +>>>>>>> master diff --git a/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 b/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 new file mode 100644 index 0000000000000000000000000000000000000000..8725b65154d298be664bde3c0b23c544c8101126 --- /dev/null +++ b/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 @@ -0,0 +1,115 @@ +{ + "train_loss_interp_naive": [ + 0.29836818281757205, + 0.2972275394162065, + 0.29598502138335453, + 0.29468635028522266, + 0.29337929414683267, + 0.2921232593883025, + 0.29098323301265117, + 0.2900340731598829, + 0.28932713030984525, + 0.2889324751143393, + 0.2888887188348331, + 0.2891541173387515, + 0.289615623750969, + 0.290175536745473, + 0.29076103964134264, + 0.29138649127593164, + 0.29207867560418027, + 0.2928481038267675, + 0.29372942124150303, + 0.29469300639864643, + 0.2957002241164446, + 0.2966967654463492, + 0.2976349547113243, + 0.2984672110331686, + 0.29918966795268814 + ], + "test_loss_interp_naive": [ + 0.29836818281757205, + 0.2972275394162065, + 0.29598502138335453, + 0.29468635028522266, + 0.29337929414683267, + 0.2921232593883025, + 0.29098323301265117, + 0.2900340731598829, + 0.28932713030984525, + 0.2889324751143393, + 0.2888887188348331, + 0.2891541173387515, + 0.289615623750969, + 0.290175536745473, + 0.29076103964134264, + 0.29138649127593164, + 0.29207867560418027, + 0.2928481038267675, + 0.29372942124150303, + 0.29469300639864643, + 0.2957002241164446, + 0.2966967654463492, + 0.2976349547113243, + 0.2984672110331686, + 0.29918966795268814 + ], + "train_acc_interp_naive": [ + 0.9067105263157895, + 0.9069736842105263, + 0.9073684210526316, + 0.9076315789473685, + 0.9075, + 0.9069736842105263, + 0.9069736842105263, + 0.9067105263157895, + 0.9063157894736842, + 0.9055263157894737, + 0.9053947368421053, + 0.9052631578947369, + 0.9053947368421053, + 0.9046052631578947, + 0.9048684210526315, + 0.9051315789473684, + 0.9048684210526315, + 0.905, + 0.9057894736842105, + 0.905921052631579, + 0.9064473684210527, + 0.9068421052631579, + 0.9073684210526316, + 0.9073684210526316, + 0.9077631578947368 + ], + "test_acc_interp_naive": [ + 0.9067105263157895, + 0.9069736842105263, + 0.9073684210526316, + 0.9076315789473685, + 0.9075, + 0.9069736842105263, + 0.9069736842105263, + 0.9067105263157895, + 0.9063157894736842, + 0.9055263157894737, + 0.9053947368421053, + 0.9052631578947369, + 0.9053947368421053, + 0.9046052631578947, + 0.9048684210526315, + 0.9051315789473684, + 0.9048684210526315, + 0.905, + 0.9057894736842105, + 0.905921052631579, + 0.9064473684210527, + 0.9068421052631579, + 0.9073684210526316, + 0.9073684210526316, + 0.9077631578947368 + ], + "train_loss_interp_clever_list": [], + "test_loss_interp_clever_list": [], + "train_acc_interp_clever_list": [], + "test_acc_interp_clever_list": [], + "baseline_train_loss": 0.29877892538513007 +} \ No newline at end of file diff --git a/results/dbpedia/[lr2e-05-epochs2-batch16-seed0+lr2e-05-epochs2-batch16-seed20].json b/results/dbpedia/[lr2e-05-epochs2-batch16-seed0+lr2e-05-epochs2-batch16-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..3b5e9eb640dd5a8db4dac770eb2046a0effeca8e --- /dev/null +++ b/results/dbpedia/[lr2e-05-epochs2-batch16-seed0+lr2e-05-epochs2-batch16-seed20].json @@ -0,0 +1,115 @@ +{ + "train_loss_interp_naive": [ + 0.21231986809444303, + 0.21220118038915098, + 0.21247703478366345, + 0.213066814562668, + 0.21389776634285226, + 0.21489918130869046, + 0.21599381488649588, + 0.2171046986388974, + 0.21814477271379695, + 0.21903549357948618, + 0.21970179501222448, + 0.22008288120426653, + 0.22013864247666465, + 0.2198468935472063, + 0.2192188657838851, + 0.21828559875467585, + 0.21710189232923505, + 0.21574412261415274, + 0.21429328908077958, + 0.21284248219735713, + 0.21148103146167058, + 0.21029916770357845, + 0.20937626931226502, + 0.20879529773857858, + 0.20863528792109962 + ], + "test_loss_interp_naive": [ + 0.20861939781067199, + 0.2087444668138706, + 0.20925309113371598, + 0.21007418510381726, + 0.2111348230448116, + 0.21236083116010407, + 0.21367340323273773, + 0.2149889662562607, + 0.2162217334881398, + 0.21728539088544307, + 0.21810557583541396, + 0.21861943222361008, + 0.21878299596779036, + 0.21857781036322999, + 0.21801239141454345, + 0.21711791628090468, + 0.21595336716410954, + 0.21459216228114097, + 0.2131237246246386, + 0.21164178506161735, + 0.2102436725077234, + 0.2090231578982282, + 0.20807068847170915, + 0.20747163428922188, + 0.20731005461667582 + ], + "train_acc_interp_naive": [ + 0.9423611111111111, + 0.94225, + 0.9424444444444444, + 0.9420277777777778, + 0.9418611111111111, + 0.9414722222222223, + 0.9410555555555555, + 0.9406388888888889, + 0.9402777777777778, + 0.9401666666666667, + 0.93975, + 0.93975, + 0.9396111111111111, + 0.9394722222222223, + 0.9395833333333333, + 0.9395555555555556, + 0.9398888888888889, + 0.9404444444444444, + 0.9411666666666667, + 0.94175, + 0.94225, + 0.9426666666666667, + 0.9427777777777778, + 0.9427222222222222, + 0.9427222222222222 + ], + "test_acc_interp_naive": [ + 0.9437279994716714, + 0.9435306115715497, + 0.943070039804599, + 0.9429384478711845, + 0.9424943250959106, + 0.942115998287344, + 0.9418199164371615, + 0.941425140636918, + 0.9408000789531992, + 0.9405368950863702, + 0.9402737112195413, + 0.9401256702944499, + 0.9401256702944499, + 0.9402572622278644, + 0.9403888541612789, + 0.9404875481113398, + 0.940816527944876, + 0.9411784057617658, + 0.9418034674454846, + 0.9423791821541729, + 0.9427410599710628, + 0.9432838766963976, + 0.9434812645965193, + 0.9432345297213671, + 0.9432509787130439 + ], + "train_loss_interp_clever_list": [], + "test_loss_interp_clever_list": [], + "train_acc_interp_clever_list": [], + "test_acc_interp_clever_list": [], + "baseline_train_loss": 0.21047757800777134 +} \ No newline at end of file diff --git a/results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json b/results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..170af0605a88b00fd2e89c3b100db580f449cf24 --- /dev/null +++ b/results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.0047, + 1.0071, + 1.0121, + 1.0204, + 1.0333, + 1.0546, + 1.0992, + 1.2213, + 1.5737, + 2.3133, + 3.2056, + 3.8128, + 3.8875, + 3.42, + 2.597, + 1.7738, + 1.2919, + 1.1168, + 1.0583, + 1.0321, + 1.0168, + 1.0068, + 1.0004, + 0.9967, + 0.9955 + ], + "Test Loss": [ + 1.0047, + 1.0071, + 1.0121, + 1.0204, + 1.0333, + 1.0546, + 1.0992, + 1.2213, + 1.5737, + 2.3133, + 3.2056, + 3.8128, + 3.8875, + 3.42, + 2.597, + 1.7738, + 1.2919, + 1.1168, + 1.0583, + 1.0321, + 1.0168, + 1.0068, + 1.0004, + 0.9967, + 0.9955 + ], + "Val PPL": [ + 2.731, + 2.7377, + 2.7515, + 2.7742, + 2.8103, + 2.8708, + 3.0017, + 3.3915, + 4.8244, + 10.1078, + 24.6708, + 45.2771, + 48.7895, + 30.5682, + 13.4233, + 5.8934, + 3.6399, + 3.055, + 2.8814, + 2.8071, + 2.7642, + 2.7368, + 2.7193, + 2.7094, + 2.7061 + ], + "Test PPL": [ + 2.731, + 2.7377, + 2.7515, + 2.7742, + 2.8103, + 2.8708, + 3.0017, + 3.3915, + 4.8244, + 10.1078, + 24.6708, + 45.2771, + 48.7895, + 30.5682, + 13.4233, + 5.8934, + 3.6399, + 3.055, + 2.8814, + 2.8071, + 2.7642, + 2.7368, + 2.7193, + 2.7094, + 2.7061 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.0047, + 1.0052, + 1.0066, + 1.0088, + 1.0117, + 1.0158, + 1.021, + 1.0277, + 1.0358, + 1.0448, + 1.0533, + 1.0584, + 1.058, + 1.0522, + 1.0432, + 1.0332, + 1.0242, + 1.0166, + 1.0105, + 1.0056, + 1.0019, + 0.9991, + 0.9971, + 0.9959, + 0.9955 + ], + "Test Loss": [ + 1.0047, + 1.0052, + 1.0066, + 1.0088, + 1.0117, + 1.0158, + 1.021, + 1.0277, + 1.0358, + 1.0448, + 1.0533, + 1.0584, + 1.058, + 1.0522, + 1.0432, + 1.0332, + 1.0242, + 1.0166, + 1.0105, + 1.0056, + 1.0019, + 0.9991, + 0.9971, + 0.9959, + 0.9955 + ], + "Val PPL": [ + 2.731, + 2.7325, + 2.7363, + 2.7422, + 2.7504, + 2.7615, + 2.776, + 2.7945, + 2.8173, + 2.843, + 2.867, + 2.8818, + 2.8806, + 2.8641, + 2.8382, + 2.8102, + 2.7848, + 2.7637, + 2.7469, + 2.7337, + 2.7234, + 2.7157, + 2.7103, + 2.7071, + 2.7061 + ], + "Test PPL": [ + 2.731, + 2.7325, + 2.7363, + 2.7422, + 2.7504, + 2.7615, + 2.776, + 2.7945, + 2.8173, + 2.843, + 2.867, + 2.8818, + 2.8806, + 2.8641, + 2.8382, + 2.8102, + 2.7848, + 2.7637, + 2.7469, + 2.7337, + 2.7234, + 2.7157, + 2.7103, + 2.7071, + 2.7061 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.0047, + 1.0052, + 1.0065, + 1.0085, + 1.0113, + 1.015, + 1.0199, + 1.0259, + 1.0334, + 1.0419, + 1.0501, + 1.0557, + 1.0564, + 1.0517, + 1.0434, + 1.0337, + 1.0247, + 1.0169, + 1.0107, + 1.0058, + 1.002, + 0.9991, + 0.9971, + 0.9959, + 0.9955 + ], + "Test Loss": [ + 1.0047, + 1.0052, + 1.0065, + 1.0085, + 1.0113, + 1.015, + 1.0199, + 1.0259, + 1.0334, + 1.0419, + 1.0501, + 1.0557, + 1.0564, + 1.0517, + 1.0434, + 1.0337, + 1.0247, + 1.0169, + 1.0107, + 1.0058, + 1.002, + 0.9991, + 0.9971, + 0.9959, + 0.9955 + ], + "Val PPL": [ + 2.731, + 2.7325, + 2.736, + 2.7415, + 2.7492, + 2.7594, + 2.7728, + 2.7897, + 2.8106, + 2.8346, + 2.8579, + 2.8739, + 2.8761, + 2.8627, + 2.8389, + 2.8115, + 2.7862, + 2.7647, + 2.7476, + 2.7341, + 2.7237, + 2.7159, + 2.7103, + 2.7071, + 2.7061 + ], + "Test PPL": [ + 2.731, + 2.7325, + 2.736, + 2.7415, + 2.7492, + 2.7594, + 2.7728, + 2.7897, + 2.8106, + 2.8346, + 2.8579, + 2.8739, + 2.8761, + 2.8627, + 2.8389, + 2.8115, + 2.7862, + 2.7647, + 2.7476, + 2.7341, + 2.7237, + 2.7159, + 2.7103, + 2.7071, + 2.7061 + ] + } +} \ No newline at end of file diff --git a/results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json b/results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..55bd2ca333fc33d880485b351b6e0e0f151689bc --- /dev/null +++ b/results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.0047, + 1.007, + 1.012, + 1.02, + 1.0326, + 1.0538, + 1.1007, + 1.2414, + 1.6943, + 2.5639, + 3.4104, + 3.919, + 3.9714, + 3.5557, + 2.7585, + 1.8631, + 1.3153, + 1.1213, + 1.0589, + 1.0321, + 1.0168, + 1.0069, + 1.0006, + 0.997, + 0.9958 + ], + "Test Loss": [ + 1.0047, + 1.007, + 1.012, + 1.02, + 1.0326, + 1.0538, + 1.1007, + 1.2414, + 1.6943, + 2.5639, + 3.4104, + 3.919, + 3.9714, + 3.5557, + 2.7585, + 1.8631, + 1.3153, + 1.1213, + 1.0589, + 1.0321, + 1.0168, + 1.0069, + 1.0006, + 0.997, + 0.9958 + ], + "Val PPL": [ + 2.731, + 2.7375, + 2.7511, + 2.7733, + 2.8084, + 2.8686, + 3.0061, + 3.4603, + 5.4427, + 12.987, + 30.2785, + 50.3487, + 53.0575, + 35.0125, + 15.7766, + 6.4437, + 3.7257, + 3.0688, + 2.8833, + 2.807, + 2.7643, + 2.7372, + 2.7198, + 2.7101, + 2.707 + ], + "Test PPL": [ + 2.731, + 2.7375, + 2.7511, + 2.7733, + 2.8084, + 2.8686, + 3.0061, + 3.4603, + 5.4427, + 12.987, + 30.2785, + 50.3487, + 53.0575, + 35.0125, + 15.7766, + 6.4437, + 3.7257, + 3.0688, + 2.8833, + 2.807, + 2.7643, + 2.7372, + 2.7198, + 2.7101, + 2.707 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.0047, + 1.0053, + 1.0067, + 1.009, + 1.0124, + 1.0168, + 1.0226, + 1.03, + 1.0392, + 1.0501, + 1.0611, + 1.0687, + 1.0692, + 1.062, + 1.0503, + 1.038, + 1.0272, + 1.0185, + 1.0117, + 1.0064, + 1.0024, + 0.9994, + 0.9973, + 0.9962, + 0.9958 + ], + "Test Loss": [ + 1.0047, + 1.0053, + 1.0067, + 1.009, + 1.0124, + 1.0168, + 1.0226, + 1.03, + 1.0392, + 1.0501, + 1.0611, + 1.0687, + 1.0692, + 1.062, + 1.0503, + 1.038, + 1.0272, + 1.0185, + 1.0117, + 1.0064, + 1.0024, + 0.9994, + 0.9973, + 0.9962, + 0.9958 + ], + "Val PPL": [ + 2.731, + 2.7326, + 2.7366, + 2.743, + 2.7521, + 2.7643, + 2.7803, + 2.801, + 2.8271, + 2.8579, + 2.8896, + 2.9115, + 2.9131, + 2.8922, + 2.8584, + 2.8235, + 2.7932, + 2.7691, + 2.7502, + 2.7357, + 2.7247, + 2.7166, + 2.711, + 2.7078, + 2.707 + ], + "Test PPL": [ + 2.731, + 2.7326, + 2.7366, + 2.743, + 2.7521, + 2.7643, + 2.7803, + 2.801, + 2.8271, + 2.8579, + 2.8896, + 2.9115, + 2.9131, + 2.8922, + 2.8584, + 2.8235, + 2.7932, + 2.7691, + 2.7502, + 2.7357, + 2.7247, + 2.7166, + 2.711, + 2.7078, + 2.707 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.0047, + 1.0053, + 1.0066, + 1.0089, + 1.012, + 1.0161, + 1.0215, + 1.0284, + 1.0371, + 1.0475, + 1.0583, + 1.0669, + 1.0693, + 1.0639, + 1.0529, + 1.0404, + 1.0291, + 1.0199, + 1.0126, + 1.007, + 1.0028, + 0.9996, + 0.9974, + 0.9962, + 0.9958 + ], + "Test Loss": [ + 1.0047, + 1.0053, + 1.0066, + 1.0089, + 1.012, + 1.0161, + 1.0215, + 1.0284, + 1.0371, + 1.0475, + 1.0583, + 1.0669, + 1.0693, + 1.0639, + 1.0529, + 1.0404, + 1.0291, + 1.0199, + 1.0126, + 1.007, + 1.0028, + 0.9996, + 0.9974, + 0.9962, + 0.9958 + ], + "Val PPL": [ + 2.731, + 2.7326, + 2.7364, + 2.7425, + 2.751, + 2.7624, + 2.7774, + 2.7966, + 2.821, + 2.8505, + 2.8815, + 2.9063, + 2.9133, + 2.8975, + 2.8659, + 2.8305, + 2.7987, + 2.773, + 2.7529, + 2.7374, + 2.7258, + 2.7172, + 2.7113, + 2.7079, + 2.707 + ], + "Test PPL": [ + 2.731, + 2.7326, + 2.7364, + 2.7425, + 2.751, + 2.7624, + 2.7774, + 2.7966, + 2.821, + 2.8505, + 2.8815, + 2.9063, + 2.9133, + 2.8975, + 2.8659, + 2.8305, + 2.7987, + 2.773, + 2.7529, + 2.7374, + 2.7258, + 2.7172, + 2.7113, + 2.7079, + 2.707 + ] + } +} \ No newline at end of file diff --git a/results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json b/results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..8cd95bd4c8eee3e1108e10e95f5990bee30f667f --- /dev/null +++ b/results/enwik8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 0.9955, + 0.9967, + 1.0002, + 1.0064, + 1.0159, + 1.0307, + 1.0565, + 1.1186, + 1.3308, + 1.9965, + 3.0682, + 3.9179, + 4.2558, + 3.9762, + 3.1691, + 2.1071, + 1.3712, + 1.1285, + 1.0606, + 1.0329, + 1.0172, + 1.0072, + 1.0007, + 0.997, + 0.9958 + ], + "Test Loss": [ + 0.9955, + 0.9967, + 1.0002, + 1.0064, + 1.0159, + 1.0307, + 1.0565, + 1.1186, + 1.3308, + 1.9965, + 3.0682, + 3.9179, + 4.2558, + 3.9762, + 3.1691, + 2.1071, + 1.3712, + 1.1285, + 1.0606, + 1.0329, + 1.0172, + 1.0072, + 1.0007, + 0.997, + 0.9958 + ], + "Val PPL": [ + 2.7061, + 2.7093, + 2.7189, + 2.7357, + 2.7619, + 2.803, + 2.8764, + 3.0607, + 3.7842, + 7.363, + 21.503, + 50.2957, + 70.514, + 53.3117, + 23.7854, + 8.2245, + 3.9402, + 3.091, + 2.8882, + 2.8092, + 2.7654, + 2.7378, + 2.7202, + 2.7103, + 2.707 + ], + "Test PPL": [ + 2.7061, + 2.7093, + 2.7189, + 2.7357, + 2.7619, + 2.803, + 2.8764, + 3.0607, + 3.7842, + 7.363, + 21.503, + 50.2957, + 70.514, + 53.3117, + 23.7854, + 8.2245, + 3.9402, + 3.091, + 2.8882, + 2.8092, + 2.7654, + 2.7378, + 2.7202, + 2.7103, + 2.707 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 0.9955, + 0.9956, + 0.9966, + 0.9984, + 1.001, + 1.0046, + 1.0096, + 1.0163, + 1.0253, + 1.0373, + 1.0514, + 1.0638, + 1.0691, + 1.0639, + 1.0513, + 1.0374, + 1.0257, + 1.0167, + 1.01, + 1.005, + 1.0013, + 0.9987, + 0.9969, + 0.996, + 0.9958 + ], + "Test Loss": [ + 0.9955, + 0.9956, + 0.9966, + 0.9984, + 1.001, + 1.0046, + 1.0096, + 1.0163, + 1.0253, + 1.0373, + 1.0514, + 1.0638, + 1.0691, + 1.0639, + 1.0513, + 1.0374, + 1.0257, + 1.0167, + 1.01, + 1.005, + 1.0013, + 0.9987, + 0.9969, + 0.996, + 0.9958 + ], + "Val PPL": [ + 2.7061, + 2.7065, + 2.709, + 2.7139, + 2.721, + 2.7309, + 2.7445, + 2.7629, + 2.788, + 2.8215, + 2.8616, + 2.8974, + 2.9127, + 2.8976, + 2.8613, + 2.8217, + 2.789, + 2.7642, + 2.7456, + 2.7319, + 2.7219, + 2.7148, + 2.7099, + 2.7073, + 2.707 + ], + "Test PPL": [ + 2.7061, + 2.7065, + 2.709, + 2.7139, + 2.721, + 2.7309, + 2.7445, + 2.7629, + 2.788, + 2.8215, + 2.8616, + 2.8974, + 2.9127, + 2.8976, + 2.8613, + 2.8217, + 2.789, + 2.7642, + 2.7456, + 2.7319, + 2.7219, + 2.7148, + 2.7099, + 2.7073, + 2.707 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 0.9955, + 0.9956, + 0.9965, + 0.9981, + 1.0005, + 1.0038, + 1.0083, + 1.0142, + 1.022, + 1.032, + 1.0443, + 1.0565, + 1.0639, + 1.0625, + 1.0527, + 1.0395, + 1.0275, + 1.018, + 1.0108, + 1.0055, + 1.0016, + 0.9988, + 0.997, + 0.996, + 0.9958 + ], + "Test Loss": [ + 0.9955, + 0.9956, + 0.9965, + 0.9981, + 1.0005, + 1.0038, + 1.0083, + 1.0142, + 1.022, + 1.032, + 1.0443, + 1.0565, + 1.0639, + 1.0625, + 1.0527, + 1.0395, + 1.0275, + 1.018, + 1.0108, + 1.0055, + 1.0016, + 0.9988, + 0.997, + 0.996, + 0.9958 + ], + "Val PPL": [ + 2.7061, + 2.7064, + 2.7088, + 2.7132, + 2.7197, + 2.7288, + 2.7409, + 2.7572, + 2.7786, + 2.8067, + 2.8413, + 2.8762, + 2.8977, + 2.8935, + 2.8655, + 2.8279, + 2.794, + 2.7676, + 2.7479, + 2.7333, + 2.7227, + 2.7151, + 2.7101, + 2.7074, + 2.707 + ], + "Test PPL": [ + 2.7061, + 2.7064, + 2.7088, + 2.7132, + 2.7197, + 2.7288, + 2.7409, + 2.7572, + 2.7786, + 2.8067, + 2.8413, + 2.8762, + 2.8977, + 2.8935, + 2.8655, + 2.8279, + 2.794, + 2.7676, + 2.7479, + 2.7333, + 2.7227, + 2.7151, + 2.7101, + 2.7074, + 2.707 + ] + } +} \ No newline at end of file diff --git a/results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json b/results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..ed75a9e8cfc1065d749bd6c2c2f5905d470b717c --- /dev/null +++ b/results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.0048, + 1.0072, + 1.0121, + 1.0201, + 1.0324, + 1.0534, + 1.1006, + 1.242, + 1.6775, + 2.5397, + 3.3603, + 3.8317, + 3.951, + 3.7364, + 3.2037, + 2.4046, + 1.6065, + 1.2036, + 1.0792, + 1.0385, + 1.0199, + 1.0088, + 1.002, + 0.9982, + 0.9972 + ], + "Test Loss": [ + 1.0048, + 1.0072, + 1.0121, + 1.0201, + 1.0324, + 1.0534, + 1.1006, + 1.242, + 1.6775, + 2.5397, + 3.3603, + 3.8317, + 3.951, + 3.7364, + 3.2037, + 2.4046, + 1.6065, + 1.2036, + 1.0792, + 1.0385, + 1.0199, + 1.0088, + 1.002, + 0.9982, + 0.9972 + ], + "Val PPL": [ + 2.7314, + 2.738, + 2.7515, + 2.7734, + 2.8078, + 2.8675, + 3.006, + 3.4625, + 5.3523, + 12.6755, + 28.7965, + 46.1401, + 51.9873, + 41.9466, + 24.6227, + 11.0745, + 4.9851, + 3.332, + 2.9422, + 2.8249, + 2.7729, + 2.7424, + 2.7236, + 2.7134, + 2.7108 + ], + "Test PPL": [ + 2.7314, + 2.738, + 2.7515, + 2.7734, + 2.8078, + 2.8675, + 3.006, + 3.4625, + 5.3523, + 12.6755, + 28.7965, + 46.1401, + 51.9873, + 41.9466, + 24.6227, + 11.0745, + 4.9851, + 3.332, + 2.9422, + 2.8249, + 2.7729, + 2.7424, + 2.7236, + 2.7134, + 2.7108 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.0048, + 1.0049, + 1.0054, + 1.0063, + 1.0076, + 1.0093, + 1.0114, + 1.0139, + 1.0168, + 1.0197, + 1.0224, + 1.0239, + 1.0236, + 1.0213, + 1.0179, + 1.014, + 1.0103, + 1.0071, + 1.0044, + 1.0021, + 1.0002, + 0.9988, + 0.9978, + 0.9973, + 0.9972 + ], + "Test Loss": [ + 1.0048, + 1.0049, + 1.0054, + 1.0063, + 1.0076, + 1.0093, + 1.0114, + 1.0139, + 1.0168, + 1.0197, + 1.0224, + 1.0239, + 1.0236, + 1.0213, + 1.0179, + 1.014, + 1.0103, + 1.0071, + 1.0044, + 1.0021, + 1.0002, + 0.9988, + 0.9978, + 0.9973, + 0.9972 + ], + "Val PPL": [ + 2.7314, + 2.7316, + 2.733, + 2.7356, + 2.7391, + 2.7437, + 2.7494, + 2.7563, + 2.7642, + 2.7725, + 2.7798, + 2.784, + 2.7831, + 2.7769, + 2.7673, + 2.7566, + 2.7465, + 2.7377, + 2.7302, + 2.724, + 2.7189, + 2.7151, + 2.7123, + 2.7109, + 2.7108 + ], + "Test PPL": [ + 2.7314, + 2.7316, + 2.733, + 2.7356, + 2.7391, + 2.7437, + 2.7494, + 2.7563, + 2.7642, + 2.7725, + 2.7798, + 2.784, + 2.7831, + 2.7769, + 2.7673, + 2.7566, + 2.7465, + 2.7377, + 2.7302, + 2.724, + 2.7189, + 2.7151, + 2.7123, + 2.7109, + 2.7108 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.0048, + 1.0049, + 1.0053, + 1.0061, + 1.0072, + 1.0086, + 1.0104, + 1.0126, + 1.0153, + 1.0182, + 1.0212, + 1.0235, + 1.0244, + 1.0232, + 1.0202, + 1.0161, + 1.012, + 1.0083, + 1.0053, + 1.0027, + 1.0006, + 0.9991, + 0.9979, + 0.9973, + 0.9972 + ], + "Test Loss": [ + 1.0048, + 1.0049, + 1.0053, + 1.0061, + 1.0072, + 1.0086, + 1.0104, + 1.0126, + 1.0153, + 1.0182, + 1.0212, + 1.0235, + 1.0244, + 1.0232, + 1.0202, + 1.0161, + 1.012, + 1.0083, + 1.0053, + 1.0027, + 1.0006, + 0.9991, + 0.9979, + 0.9973, + 0.9972 + ], + "Val PPL": [ + 2.7314, + 2.7315, + 2.7327, + 2.7348, + 2.7379, + 2.7418, + 2.7468, + 2.7529, + 2.7601, + 2.7682, + 2.7764, + 2.783, + 2.7855, + 2.7821, + 2.7736, + 2.7623, + 2.7511, + 2.7411, + 2.7326, + 2.7256, + 2.72, + 2.7157, + 2.7126, + 2.711, + 2.7108 + ], + "Test PPL": [ + 2.7314, + 2.7315, + 2.7327, + 2.7348, + 2.7379, + 2.7418, + 2.7468, + 2.7529, + 2.7601, + 2.7682, + 2.7764, + 2.783, + 2.7855, + 2.7821, + 2.7736, + 2.7623, + 2.7511, + 2.7411, + 2.7326, + 2.7256, + 2.72, + 2.7157, + 2.7126, + 2.711, + 2.7108 + ] + } +} \ No newline at end of file diff --git a/results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json b/results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..3ba37ee73b5c54f32070d509039ba5048470359f --- /dev/null +++ b/results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.0048, + 1.0071, + 1.012, + 1.0198, + 1.0319, + 1.0517, + 1.0908, + 1.1923, + 1.4926, + 2.1376, + 2.8529, + 3.3063, + 3.4331, + 3.2546, + 2.8017, + 2.1408, + 1.485, + 1.1778, + 1.0789, + 1.0405, + 1.0211, + 1.0095, + 1.0024, + 0.9986, + 0.9978 + ], + "Test Loss": [ + 1.0048, + 1.0071, + 1.012, + 1.0198, + 1.0319, + 1.0517, + 1.0908, + 1.1923, + 1.4926, + 2.1376, + 2.8529, + 3.3063, + 3.4331, + 3.2546, + 2.8017, + 2.1408, + 1.485, + 1.1778, + 1.0789, + 1.0405, + 1.0211, + 1.0095, + 1.0024, + 0.9986, + 0.9978 + ], + "Val PPL": [ + 2.7314, + 2.7378, + 2.751, + 2.7726, + 2.8064, + 2.8625, + 2.9766, + 3.2945, + 4.4486, + 8.4792, + 17.3372, + 27.2838, + 30.9727, + 25.9082, + 16.4722, + 8.5058, + 4.4149, + 3.2473, + 2.9416, + 2.8306, + 2.7762, + 2.7442, + 2.7247, + 2.7145, + 2.7123 + ], + "Test PPL": [ + 2.7314, + 2.7378, + 2.751, + 2.7726, + 2.8064, + 2.8625, + 2.9766, + 3.2945, + 4.4486, + 8.4792, + 17.3372, + 27.2838, + 30.9727, + 25.9082, + 16.4722, + 8.5058, + 4.4149, + 3.2473, + 2.9416, + 2.8306, + 2.7762, + 2.7442, + 2.7247, + 2.7145, + 2.7123 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.0048, + 1.0048, + 1.0053, + 1.0063, + 1.0076, + 1.0093, + 1.0116, + 1.0147, + 1.0193, + 1.0286, + 1.0526, + 1.1092, + 1.1843, + 1.2349, + 1.2427, + 1.1989, + 1.1011, + 1.0319, + 1.0108, + 1.0043, + 1.0012, + 0.9993, + 0.9982, + 0.9977, + 0.9978 + ], + "Test Loss": [ + 1.0048, + 1.0048, + 1.0053, + 1.0063, + 1.0076, + 1.0093, + 1.0116, + 1.0147, + 1.0193, + 1.0286, + 1.0526, + 1.1092, + 1.1843, + 1.2349, + 1.2427, + 1.1989, + 1.1011, + 1.0319, + 1.0108, + 1.0043, + 1.0012, + 0.9993, + 0.9982, + 0.9977, + 0.9978 + ], + "Val PPL": [ + 2.7314, + 2.7315, + 2.7329, + 2.7354, + 2.7389, + 2.7437, + 2.75, + 2.7584, + 2.7713, + 2.7971, + 2.865, + 3.032, + 3.2684, + 3.438, + 3.4649, + 3.3165, + 3.0073, + 2.8063, + 2.7478, + 2.73, + 2.7215, + 2.7165, + 2.7134, + 2.7121, + 2.7124 + ], + "Test PPL": [ + 2.7314, + 2.7315, + 2.7329, + 2.7354, + 2.7389, + 2.7437, + 2.75, + 2.7584, + 2.7713, + 2.7971, + 2.865, + 3.032, + 3.2684, + 3.438, + 3.4649, + 3.3165, + 3.0073, + 2.8063, + 2.7478, + 2.73, + 2.7215, + 2.7165, + 2.7134, + 2.7121, + 2.7124 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.0048, + 1.0048, + 1.0052, + 1.006, + 1.0071, + 1.0086, + 1.0106, + 1.0131, + 1.0166, + 1.022, + 1.0333, + 1.0611, + 1.1194, + 1.1875, + 1.2265, + 1.2204, + 1.1553, + 1.0609, + 1.0176, + 1.0061, + 1.0019, + 0.9997, + 0.9983, + 0.9977, + 0.9978 + ], + "Test Loss": [ + 1.0048, + 1.0048, + 1.0052, + 1.006, + 1.0071, + 1.0086, + 1.0106, + 1.0131, + 1.0166, + 1.022, + 1.0333, + 1.0611, + 1.1194, + 1.1875, + 1.2265, + 1.2204, + 1.1553, + 1.0609, + 1.0176, + 1.0061, + 1.0019, + 0.9997, + 0.9983, + 0.9977, + 0.9978 + ], + "Val PPL": [ + 2.7314, + 2.7314, + 2.7325, + 2.7347, + 2.7378, + 2.7418, + 2.7472, + 2.7541, + 2.7637, + 2.7789, + 2.8102, + 2.8894, + 3.0629, + 3.2787, + 3.4094, + 3.3886, + 3.1749, + 2.8888, + 2.7666, + 2.7351, + 2.7235, + 2.7174, + 2.7138, + 2.7121, + 2.7123 + ], + "Test PPL": [ + 2.7314, + 2.7314, + 2.7325, + 2.7347, + 2.7378, + 2.7418, + 2.7472, + 2.7541, + 2.7637, + 2.7789, + 2.8102, + 2.8894, + 3.0629, + 3.2787, + 3.4094, + 3.3886, + 3.1749, + 2.8888, + 2.7666, + 2.7351, + 2.7235, + 2.7174, + 2.7138, + 2.7121, + 2.7123 + ] + } +} \ No newline at end of file diff --git a/results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json b/results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..af7d24c9d97fe6d8ae74589ee9b4478210c323ec --- /dev/null +++ b/results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 0.9972, + 0.9981, + 1.0018, + 1.0084, + 1.019, + 1.037, + 1.0785, + 1.2137, + 1.6466, + 2.5233, + 3.3907, + 3.9375, + 4.1491, + 4.0052, + 3.5155, + 2.6962, + 1.7777, + 1.2681, + 1.0977, + 1.0438, + 1.0216, + 1.0096, + 1.0024, + 0.9986, + 0.9978 + ], + "Test Loss": [ + 0.9972, + 0.9981, + 1.0018, + 1.0084, + 1.019, + 1.037, + 1.0785, + 1.2137, + 1.6466, + 2.5233, + 3.3907, + 3.9375, + 4.1491, + 4.0052, + 3.5155, + 2.6962, + 1.7777, + 1.2681, + 1.0977, + 1.0438, + 1.0216, + 1.0096, + 1.0024, + 0.9986, + 0.9978 + ], + "Val PPL": [ + 2.7108, + 2.7132, + 2.7231, + 2.7412, + 2.7704, + 2.8208, + 2.9402, + 3.3659, + 5.1894, + 12.4694, + 29.6863, + 51.2913, + 63.3755, + 54.8847, + 33.6314, + 14.8239, + 5.916, + 3.5542, + 2.9973, + 2.84, + 2.7777, + 2.7444, + 2.7247, + 2.7145, + 2.7123 + ], + "Test PPL": [ + 2.7108, + 2.7132, + 2.7231, + 2.7412, + 2.7704, + 2.8208, + 2.9402, + 3.3659, + 5.1894, + 12.4694, + 29.6863, + 51.2913, + 63.3755, + 54.8847, + 33.6314, + 14.8239, + 5.916, + 3.5542, + 2.9973, + 2.84, + 2.7777, + 2.7444, + 2.7247, + 2.7145, + 2.7123 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 0.9972, + 0.9971, + 0.9974, + 0.9982, + 0.9995, + 1.0013, + 1.0039, + 1.0073, + 1.0125, + 1.022, + 1.0447, + 1.1004, + 1.1772, + 1.2196, + 1.2006, + 1.1256, + 1.0513, + 1.0187, + 1.0077, + 1.003, + 1.0004, + 0.9987, + 0.9978, + 0.9975, + 0.9978 + ], + "Test Loss": [ + 0.9972, + 0.9971, + 0.9974, + 0.9982, + 0.9995, + 1.0013, + 1.0039, + 1.0073, + 1.0125, + 1.022, + 1.0447, + 1.1004, + 1.1772, + 1.2196, + 1.2006, + 1.1256, + 1.0513, + 1.0187, + 1.0077, + 1.003, + 1.0004, + 0.9987, + 0.9978, + 0.9975, + 0.9978 + ], + "Val PPL": [ + 2.7108, + 2.7104, + 2.7112, + 2.7134, + 2.7169, + 2.7219, + 2.7288, + 2.7383, + 2.7525, + 2.7786, + 2.8425, + 3.0055, + 3.2451, + 3.3857, + 3.3221, + 3.0822, + 2.8615, + 2.7695, + 2.7393, + 2.7264, + 2.7193, + 2.7148, + 2.7124, + 2.7115, + 2.7124 + ], + "Test PPL": [ + 2.7108, + 2.7104, + 2.7112, + 2.7134, + 2.7169, + 2.7219, + 2.7288, + 2.7383, + 2.7525, + 2.7786, + 2.8425, + 3.0055, + 3.2451, + 3.3857, + 3.3221, + 3.0822, + 2.8615, + 2.7695, + 2.7393, + 2.7264, + 2.7193, + 2.7148, + 2.7124, + 2.7115, + 2.7124 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 0.9972, + 0.997, + 0.9973, + 0.9979, + 0.999, + 1.0005, + 1.0025, + 1.0052, + 1.0088, + 1.0143, + 1.0244, + 1.0493, + 1.1047, + 1.1736, + 1.2015, + 1.1634, + 1.0826, + 1.0291, + 1.0108, + 1.0042, + 1.0009, + 0.999, + 0.9979, + 0.9975, + 0.9978 + ], + "Test Loss": [ + 0.9972, + 0.997, + 0.9973, + 0.9979, + 0.999, + 1.0005, + 1.0025, + 1.0052, + 1.0088, + 1.0143, + 1.0244, + 1.0493, + 1.1047, + 1.1736, + 1.2015, + 1.1634, + 1.0826, + 1.0291, + 1.0108, + 1.0042, + 1.0009, + 0.999, + 0.9979, + 0.9975, + 0.9978 + ], + "Val PPL": [ + 2.7108, + 2.7103, + 2.711, + 2.7127, + 2.7155, + 2.7196, + 2.7251, + 2.7324, + 2.7424, + 2.7573, + 2.7855, + 2.8556, + 3.0184, + 3.2335, + 3.3251, + 3.2007, + 2.9524, + 2.7985, + 2.7477, + 2.7297, + 2.7207, + 2.7155, + 2.7126, + 2.7115, + 2.7123 + ], + "Test PPL": [ + 2.7108, + 2.7103, + 2.711, + 2.7127, + 2.7155, + 2.7196, + 2.7251, + 2.7324, + 2.7424, + 2.7573, + 2.7855, + 2.8556, + 3.0184, + 3.2335, + 3.3251, + 3.2007, + 2.9524, + 2.7985, + 2.7477, + 2.7297, + 2.7207, + 2.7155, + 2.7126, + 2.7115, + 2.7123 + ] + } +} \ No newline at end of file diff --git a/results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json b/results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..297424fdf1f80c9f5f0bd35523be7d7ef2c99750 --- /dev/null +++ b/results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.0051, + 1.0071, + 1.0119, + 1.02, + 1.0331, + 1.0553, + 1.1027, + 1.2476, + 1.7547, + 2.7906, + 3.8529, + 4.5533, + 4.6814, + 4.2446, + 3.4481, + 2.4965, + 1.6157, + 1.1988, + 1.0801, + 1.0407, + 1.021, + 1.0092, + 1.0021, + 0.9982, + 0.9974 + ], + "Test Loss": [ + 1.0051, + 1.0071, + 1.0119, + 1.02, + 1.0331, + 1.0553, + 1.1027, + 1.2476, + 1.7547, + 2.7906, + 3.8529, + 4.5533, + 4.6814, + 4.2446, + 3.4481, + 2.4965, + 1.6157, + 1.1988, + 1.0801, + 1.0407, + 1.021, + 1.0092, + 1.0021, + 0.9982, + 0.9974 + ], + "Val PPL": [ + 2.7321, + 2.7375, + 2.7508, + 2.7732, + 2.8098, + 2.8729, + 3.0124, + 3.4818, + 5.7819, + 16.2901, + 47.1272, + 94.9449, + 107.9215, + 69.7294, + 31.4406, + 12.1401, + 5.0314, + 3.3162, + 2.945, + 2.8311, + 2.776, + 2.7435, + 2.7239, + 2.7134, + 2.7111 + ], + "Test PPL": [ + 2.7321, + 2.7375, + 2.7508, + 2.7732, + 2.8098, + 2.8729, + 3.0124, + 3.4818, + 5.7819, + 16.2901, + 47.1272, + 94.9449, + 107.9215, + 69.7294, + 31.4406, + 12.1401, + 5.0314, + 3.3162, + 2.945, + 2.8311, + 2.776, + 2.7435, + 2.7239, + 2.7134, + 2.7111 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.0051, + 1.0052, + 1.0061, + 1.0078, + 1.0104, + 1.0141, + 1.0194, + 1.0268, + 1.0374, + 1.0522, + 1.0701, + 1.0852, + 1.0884, + 1.0774, + 1.0588, + 1.0413, + 1.0278, + 1.0181, + 1.0112, + 1.0062, + 1.0026, + 1.0, + 0.9983, + 0.9974, + 0.9974 + ], + "Test Loss": [ + 1.0051, + 1.0052, + 1.0061, + 1.0078, + 1.0104, + 1.0141, + 1.0194, + 1.0268, + 1.0374, + 1.0522, + 1.0701, + 1.0852, + 1.0884, + 1.0774, + 1.0588, + 1.0413, + 1.0278, + 1.0181, + 1.0112, + 1.0062, + 1.0026, + 1.0, + 0.9983, + 0.9974, + 0.9974 + ], + "Val PPL": [ + 2.7321, + 2.7323, + 2.7348, + 2.7395, + 2.7467, + 2.757, + 2.7714, + 2.792, + 2.8218, + 2.8641, + 2.9158, + 2.9602, + 2.9695, + 2.937, + 2.8829, + 2.8329, + 2.7949, + 2.768, + 2.749, + 2.7352, + 2.7252, + 2.7182, + 2.7136, + 2.7113, + 2.7112 + ], + "Test PPL": [ + 2.7321, + 2.7323, + 2.7348, + 2.7395, + 2.7467, + 2.757, + 2.7714, + 2.792, + 2.8218, + 2.8641, + 2.9158, + 2.9602, + 2.9695, + 2.937, + 2.8829, + 2.8329, + 2.7949, + 2.768, + 2.749, + 2.7352, + 2.7252, + 2.7182, + 2.7136, + 2.7113, + 2.7112 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.0051, + 1.0051, + 1.006, + 1.0075, + 1.0098, + 1.0132, + 1.0178, + 1.0243, + 1.0337, + 1.0473, + 1.0655, + 1.0844, + 1.0939, + 1.0871, + 1.0683, + 1.0479, + 1.0319, + 1.0206, + 1.0127, + 1.0071, + 1.0031, + 1.0003, + 0.9984, + 0.9975, + 0.9974 + ], + "Test Loss": [ + 1.0051, + 1.0051, + 1.006, + 1.0075, + 1.0098, + 1.0132, + 1.0178, + 1.0243, + 1.0337, + 1.0473, + 1.0655, + 1.0844, + 1.0939, + 1.0871, + 1.0683, + 1.0479, + 1.0319, + 1.0206, + 1.0127, + 1.0071, + 1.0031, + 1.0003, + 0.9984, + 0.9975, + 0.9974 + ], + "Val PPL": [ + 2.7321, + 2.7323, + 2.7346, + 2.7388, + 2.7452, + 2.7543, + 2.767, + 2.7851, + 2.8115, + 2.8499, + 2.9022, + 2.9577, + 2.986, + 2.9655, + 2.9104, + 2.8518, + 2.8063, + 2.7749, + 2.7531, + 2.7377, + 2.7267, + 2.719, + 2.714, + 2.7114, + 2.7112 + ], + "Test PPL": [ + 2.7321, + 2.7323, + 2.7346, + 2.7388, + 2.7452, + 2.7543, + 2.767, + 2.7851, + 2.8115, + 2.8499, + 2.9022, + 2.9577, + 2.986, + 2.9655, + 2.9104, + 2.8518, + 2.8063, + 2.7749, + 2.7531, + 2.7377, + 2.7267, + 2.719, + 2.714, + 2.7114, + 2.7112 + ] + } +} \ No newline at end of file diff --git a/results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json b/results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..db044326e8cc50eb6cb42e848f5178c5c418855b --- /dev/null +++ b/results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.0051, + 1.007, + 1.0117, + 1.0196, + 1.0321, + 1.0538, + 1.1022, + 1.2507, + 1.6776, + 2.433, + 3.2565, + 3.8439, + 3.9582, + 3.5562, + 2.7774, + 1.9406, + 1.3874, + 1.1513, + 1.0686, + 1.0356, + 1.0183, + 1.0076, + 1.0008, + 0.997, + 0.9958 + ], + "Test Loss": [ + 1.0051, + 1.007, + 1.0117, + 1.0196, + 1.0321, + 1.0538, + 1.1022, + 1.2507, + 1.6776, + 2.433, + 3.2565, + 3.8439, + 3.9582, + 3.5562, + 2.7774, + 1.9406, + 1.3874, + 1.1513, + 1.0686, + 1.0356, + 1.0183, + 1.0076, + 1.0008, + 0.997, + 0.9958 + ], + "Val PPL": [ + 2.7321, + 2.7375, + 2.7504, + 2.7721, + 2.8071, + 2.8686, + 3.0109, + 3.4929, + 5.3527, + 11.3925, + 25.9587, + 46.7067, + 52.3641, + 35.0292, + 16.0764, + 6.963, + 4.0043, + 3.1623, + 2.9114, + 2.8168, + 2.7684, + 2.739, + 2.7205, + 2.7101, + 2.7069 + ], + "Test PPL": [ + 2.7321, + 2.7375, + 2.7504, + 2.7721, + 2.8071, + 2.8686, + 3.0109, + 3.4929, + 5.3527, + 11.3925, + 25.9587, + 46.7067, + 52.3641, + 35.0292, + 16.0764, + 6.963, + 4.0043, + 3.1623, + 2.9114, + 2.8168, + 2.7684, + 2.739, + 2.7205, + 2.7101, + 2.7069 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.0051, + 1.0052, + 1.0061, + 1.0077, + 1.01, + 1.0132, + 1.0173, + 1.0227, + 1.0292, + 1.0368, + 1.0443, + 1.0495, + 1.0501, + 1.0456, + 1.0378, + 1.029, + 1.021, + 1.0144, + 1.009, + 1.0047, + 1.0014, + 0.9989, + 0.9971, + 0.9961, + 0.9958 + ], + "Test Loss": [ + 1.0051, + 1.0052, + 1.0061, + 1.0077, + 1.01, + 1.0132, + 1.0173, + 1.0227, + 1.0292, + 1.0368, + 1.0443, + 1.0495, + 1.0501, + 1.0456, + 1.0378, + 1.029, + 1.021, + 1.0144, + 1.009, + 1.0047, + 1.0014, + 0.9989, + 0.9971, + 0.9961, + 0.9958 + ], + "Val PPL": [ + 2.7321, + 2.7324, + 2.7348, + 2.7392, + 2.7456, + 2.7543, + 2.7658, + 2.7806, + 2.7989, + 2.8201, + 2.8413, + 2.8562, + 2.858, + 2.8452, + 2.8231, + 2.7984, + 2.776, + 2.7576, + 2.7428, + 2.7312, + 2.7221, + 2.7154, + 2.7105, + 2.7078, + 2.7069 + ], + "Test PPL": [ + 2.7321, + 2.7324, + 2.7348, + 2.7392, + 2.7456, + 2.7543, + 2.7658, + 2.7806, + 2.7989, + 2.8201, + 2.8413, + 2.8562, + 2.858, + 2.8452, + 2.8231, + 2.7984, + 2.776, + 2.7576, + 2.7428, + 2.7312, + 2.7221, + 2.7154, + 2.7105, + 2.7078, + 2.7069 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.0051, + 1.0052, + 1.006, + 1.0075, + 1.0097, + 1.0126, + 1.0165, + 1.0215, + 1.0278, + 1.0351, + 1.0427, + 1.0489, + 1.0511, + 1.0481, + 1.0408, + 1.0319, + 1.0233, + 1.016, + 1.0101, + 1.0055, + 1.0019, + 0.9992, + 0.9973, + 0.9961, + 0.9958 + ], + "Test Loss": [ + 1.0051, + 1.0052, + 1.006, + 1.0075, + 1.0097, + 1.0126, + 1.0165, + 1.0215, + 1.0278, + 1.0351, + 1.0427, + 1.0489, + 1.0511, + 1.0481, + 1.0408, + 1.0319, + 1.0233, + 1.016, + 1.0101, + 1.0055, + 1.0019, + 0.9992, + 0.9973, + 0.9961, + 0.9958 + ], + "Val PPL": [ + 2.7321, + 2.7323, + 2.7347, + 2.7387, + 2.7447, + 2.7528, + 2.7636, + 2.7775, + 2.7949, + 2.8154, + 2.837, + 2.8544, + 2.8609, + 2.8522, + 2.8315, + 2.8064, + 2.7824, + 2.7622, + 2.746, + 2.7332, + 2.7234, + 2.716, + 2.7109, + 2.7078, + 2.7069 + ], + "Test PPL": [ + 2.7321, + 2.7323, + 2.7347, + 2.7387, + 2.7447, + 2.7528, + 2.7636, + 2.7775, + 2.7949, + 2.8154, + 2.837, + 2.8544, + 2.8609, + 2.8522, + 2.8315, + 2.8064, + 2.7824, + 2.7622, + 2.746, + 2.7332, + 2.7234, + 2.716, + 2.7109, + 2.7078, + 2.7069 + ] + } +} \ No newline at end of file diff --git a/results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json b/results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..e2955ecb5ed1e0dd63f3f32c167eb110c41a4f29 --- /dev/null +++ b/results/enwik8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 0.9974, + 0.9981, + 1.0016, + 1.0082, + 1.0188, + 1.0367, + 1.0731, + 1.1835, + 1.5632, + 2.35, + 3.2535, + 3.9713, + 4.2475, + 3.9425, + 3.1482, + 2.1402, + 1.426, + 1.1494, + 1.0662, + 1.0345, + 1.0176, + 1.0073, + 1.0007, + 0.9969, + 0.9958 + ], + "Test Loss": [ + 0.9974, + 0.9981, + 1.0016, + 1.0082, + 1.0188, + 1.0367, + 1.0731, + 1.1835, + 1.5632, + 2.35, + 3.2535, + 3.9713, + 4.2475, + 3.9425, + 3.1482, + 2.1402, + 1.426, + 1.1494, + 1.0662, + 1.0345, + 1.0176, + 1.0073, + 1.0007, + 0.9969, + 0.9958 + ], + "Val PPL": [ + 2.7111, + 2.7132, + 2.7227, + 2.7407, + 2.77, + 2.8198, + 2.9245, + 3.2658, + 4.774, + 10.4853, + 25.8803, + 53.0549, + 69.9269, + 51.5461, + 23.2936, + 8.5011, + 4.1621, + 3.1564, + 2.9043, + 2.8137, + 2.7667, + 2.7381, + 2.7201, + 2.7099, + 2.7069 + ], + "Test PPL": [ + 2.7111, + 2.7132, + 2.7227, + 2.7407, + 2.77, + 2.8198, + 2.9245, + 3.2658, + 4.774, + 10.4853, + 25.8803, + 53.0549, + 69.9269, + 51.5461, + 23.2936, + 8.5011, + 4.1621, + 3.1564, + 2.9043, + 2.8137, + 2.7667, + 2.7381, + 2.7201, + 2.7099, + 2.7069 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 0.9974, + 0.9972, + 0.9976, + 0.9987, + 1.0005, + 1.003, + 1.0062, + 1.0103, + 1.0152, + 1.0208, + 1.0261, + 1.03, + 1.031, + 1.0289, + 1.0243, + 1.0188, + 1.0134, + 1.0086, + 1.0048, + 1.0017, + 0.9993, + 0.9976, + 0.9964, + 0.9958, + 0.9958 + ], + "Test Loss": [ + 0.9974, + 0.9972, + 0.9976, + 0.9987, + 1.0005, + 1.003, + 1.0062, + 1.0103, + 1.0152, + 1.0208, + 1.0261, + 1.03, + 1.031, + 1.0289, + 1.0243, + 1.0188, + 1.0134, + 1.0086, + 1.0048, + 1.0017, + 0.9993, + 0.9976, + 0.9964, + 0.9958, + 0.9958 + ], + "Val PPL": [ + 2.7111, + 2.7106, + 2.7119, + 2.7149, + 2.7196, + 2.7264, + 2.7352, + 2.7464, + 2.76, + 2.7753, + 2.7903, + 2.8011, + 2.804, + 2.7979, + 2.7852, + 2.7698, + 2.7548, + 2.7418, + 2.7313, + 2.7229, + 2.7164, + 2.7117, + 2.7084, + 2.7068, + 2.7069 + ], + "Test PPL": [ + 2.7111, + 2.7106, + 2.7119, + 2.7149, + 2.7196, + 2.7264, + 2.7352, + 2.7464, + 2.76, + 2.7753, + 2.7903, + 2.8011, + 2.804, + 2.7979, + 2.7852, + 2.7698, + 2.7548, + 2.7418, + 2.7313, + 2.7229, + 2.7164, + 2.7117, + 2.7084, + 2.7068, + 2.7069 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 0.9974, + 0.9972, + 0.9975, + 0.9985, + 1.0001, + 1.0023, + 1.0052, + 1.009, + 1.0135, + 1.0187, + 1.0242, + 1.0287, + 1.031, + 1.0299, + 1.0261, + 1.0206, + 1.0149, + 1.0098, + 1.0056, + 1.0022, + 0.9996, + 0.9977, + 0.9964, + 0.9958, + 0.9958 + ], + "Test Loss": [ + 0.9974, + 0.9972, + 0.9975, + 0.9985, + 1.0001, + 1.0023, + 1.0052, + 1.009, + 1.0135, + 1.0187, + 1.0242, + 1.0287, + 1.031, + 1.0299, + 1.0261, + 1.0206, + 1.0149, + 1.0098, + 1.0056, + 1.0022, + 0.9996, + 0.9977, + 0.9964, + 0.9958, + 0.9958 + ], + "Val PPL": [ + 2.7111, + 2.7106, + 2.7116, + 2.7143, + 2.7186, + 2.7246, + 2.7326, + 2.7427, + 2.7552, + 2.7697, + 2.7849, + 2.7975, + 2.8037, + 2.8009, + 2.7901, + 2.775, + 2.7591, + 2.7451, + 2.7335, + 2.7242, + 2.7172, + 2.712, + 2.7086, + 2.7068, + 2.7069 + ], + "Test PPL": [ + 2.7111, + 2.7106, + 2.7116, + 2.7143, + 2.7186, + 2.7246, + 2.7326, + 2.7427, + 2.7552, + 2.7697, + 2.7849, + 2.7975, + 2.8037, + 2.8009, + 2.7901, + 2.775, + 2.7591, + 2.7451, + 2.7335, + 2.7242, + 2.7172, + 2.712, + 2.7086, + 2.7068, + 2.7069 + ] + } +} \ No newline at end of file diff --git a/results/imagenet/finetune-learnable-indice0-heads1-shared1-routed0-topk0-mlpFalse-seed0+finetune-learnable-indice0-heads1-shared1-routed0-topk0-mlpFalse-seed20.json b/results/imagenet/finetune-learnable-indice0-heads1-shared1-routed0-topk0-mlpFalse-seed0+finetune-learnable-indice0-heads1-shared1-routed0-topk0-mlpFalse-seed20.json new file mode 100644 index 0000000000000000000000000000000000000000..b52cd746b1deac1a355274075ca604abc8f4f213 --- /dev/null +++ b/results/imagenet/finetune-learnable-indice0-heads1-shared1-routed0-topk0-mlpFalse-seed0+finetune-learnable-indice0-heads1-shared1-routed0-topk0-mlpFalse-seed20.json @@ -0,0 +1,26 @@ +{ + "Naive": { + "Test Loss": [ + 2.75, + 7.8438, + 2.9688 + ], + "Test Acc": [ + 32.4219, + 9.375, + 33.5938 + ] + }, + "permu_head_init_ortho_opt": { + "Test Loss": [ + 2.7969, + 3.5, + 2.9219 + ], + "Test Acc": [ + 31.25, + 25.7812, + 32.4219 + ] + } +} \ No newline at end of file diff --git a/results/imagenet/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20.json b/results/imagenet/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20.json new file mode 100644 index 0000000000000000000000000000000000000000..3e7fe73ca3cc7aeaa626d8c141bf309ae98e7942 --- /dev/null +++ b/results/imagenet/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20.json @@ -0,0 +1,170 @@ +{ + "Naive": { + "Test Loss": [ + 0.6719, + 0.6758, + 0.668, + 0.6719, + 0.668, + 0.6719, + 0.668, + 0.6719, + 0.668, + 0.668, + 0.6719, + 0.668, + 0.668, + 0.668, + 0.6641, + 0.6641, + 0.6641, + 0.6641, + 0.6641, + 0.6641, + 0.6602, + 0.6602, + 0.6602, + 0.6602, + 0.6602 + ], + "Test Acc": [ + 84.7656, + 84.375, + 84.7656, + 84.375, + 83.9844, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 83.9844, + 84.375, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844 + ] + }, + "permu_head_init_ortho_no_opt": { + "Test Loss": [ + 0.6719, + 0.6758, + 0.6719, + 0.6719, + 0.6719, + 0.668, + 0.6719, + 0.6719, + 0.668, + 0.668, + 0.668, + 0.668, + 0.668, + 0.668, + 0.668, + 0.6641, + 0.6641, + 0.6641, + 0.6641, + 0.6602, + 0.6641, + 0.6602, + 0.6602, + 0.6602, + 0.6602 + ], + "Test Acc": [ + 84.7656, + 84.375, + 84.375, + 84.7656, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844 + ] + }, + "permu_head_init_ortho_opt": { + "Test Loss": [ + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.668, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.668, + 0.668, + 0.668, + 0.668, + 0.668, + 0.668, + 0.6641, + 0.6641, + 0.6641, + 0.6641, + 0.6641, + 0.6641, + 0.6602, + 0.6602, + 0.6602 + ], + "Test Acc": [ + 84.7656, + 84.375, + 84.375, + 84.7656, + 84.7656, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844, + 83.9844 + ] + } +} \ No newline at end of file diff --git a/results/imagenet/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40.json b/results/imagenet/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40.json new file mode 100644 index 0000000000000000000000000000000000000000..8685d4552866f6ecb455d2f7330f69fe4cd8848b --- /dev/null +++ b/results/imagenet/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40.json @@ -0,0 +1,170 @@ +{ + "Naive": { + "Test Loss": [ + 0.6719, + 0.6758, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.668, + 0.6719, + 0.6719, + 0.668, + 0.668, + 0.668, + 0.668, + 0.6641, + 0.668, + 0.668, + 0.6641, + 0.6641, + 0.6641 + ], + "Test Acc": [ + 84.7656, + 83.9844, + 84.7656, + 84.7656, + 84.7656, + 84.375, + 84.375, + 84.375, + 84.375, + 84.7656, + 83.9844, + 84.375, + 84.375, + 84.375, + 83.9844, + 84.375, + 84.375, + 84.375, + 84.375, + 83.9844, + 83.9844, + 83.9844, + 83.5938, + 83.5938, + 84.375 + ] + }, + "permu_head_init_ortho_no_opt": { + "Test Loss": [ + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6758, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.668, + 0.6719, + 0.668, + 0.668, + 0.668, + 0.668, + 0.6641, + 0.668, + 0.6641, + 0.6641 + ], + "Test Acc": [ + 84.7656, + 84.375, + 84.375, + 84.7656, + 84.375, + 84.375, + 84.7656, + 84.375, + 84.375, + 84.375, + 84.7656, + 84.375, + 83.9844, + 83.9844, + 83.9844, + 84.7656, + 84.375, + 83.9844, + 83.9844, + 83.9844, + 84.375, + 83.9844, + 83.5938, + 83.9844, + 83.9844 + ] + }, + "permu_head_init_ortho_opt": { + "Test Loss": [ + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6758, + 0.6758, + 0.6758, + 0.6758, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.6758, + 0.6719, + 0.6719, + 0.6719, + 0.6719, + 0.668, + 0.668, + 0.6641, + 0.6641, + 0.6641, + 0.6602, + 0.6641 + ], + "Test Acc": [ + 84.7656, + 84.375, + 84.375, + 84.7656, + 84.7656, + 84.7656, + 84.7656, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 84.375, + 83.9844, + 83.9844, + 83.5938, + 83.5938 + ] + } +} \ No newline at end of file diff --git a/results/imagenet/finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20.json b/results/imagenet/finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20.json new file mode 100644 index 0000000000000000000000000000000000000000..9c3a5552d73b2ffcec5d7c75de3209e0c195c52e --- /dev/null +++ b/results/imagenet/finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20.json @@ -0,0 +1,26 @@ +{ + "Naive": { + "Test Loss": [ + 1.2188, + 6.875, + 1.1875 + ], + "Test Acc": [ + 67.1875, + 10.1562, + 65.2344 + ] + }, + "permu_head_init_ortho_opt": { + "Test Loss": [ + 0.9219, + 2.3125, + 1.2266 + ], + "Test Acc": [ + 75.0, + 46.4844, + 65.625 + ] + } +} \ No newline at end of file diff --git a/results/imagenet/finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.json b/results/imagenet/finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.json new file mode 100644 index 0000000000000000000000000000000000000000..ff5f1795f223080f68b0dc2f73033b8ab5d23b63 --- /dev/null +++ b/results/imagenet/finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.json @@ -0,0 +1,26 @@ +{ + "Naive": { + "Test Loss": [ + 2.5, + 7.1875, + 2.5312 + ], + "Test Acc": [ + 42.1875, + 7.8125, + 41.7969 + ] + }, + "permu_head_init_ortho_opt": { + "Test Loss": [ + 2.5938, + 2.8594, + 2.6406 + ], + "Test Acc": [ + 41.4062, + 36.7188, + 43.3594 + ] + } +} \ No newline at end of file diff --git a/results/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed0+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.json b/results/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed0+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.json new file mode 100644 index 0000000000000000000000000000000000000000..7b6428cce6408c4d25e12db91d06aa4e18be0c87 --- /dev/null +++ b/results/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed0+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.json @@ -0,0 +1,26 @@ +{ + "Naive": { + "Test Loss": [ + 0.9023, + 1.9375, + 0.9883 + ], + "Test Acc": [ + 69.5312, + 41.4062, + 68.75 + ] + }, + "permu_head_init_ortho_opt": { + "Test Loss": [ + 0.9102, + 0.8789, + 0.9141 + ], + "Test Acc": [ + 70.3125, + 69.5312, + 70.7031 + ] + } +} \ No newline at end of file diff --git a/results/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.json b/results/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.json new file mode 100644 index 0000000000000000000000000000000000000000..ccde8ccff24d720b0413082481885f46e6d19c3c --- /dev/null +++ b/results/imagenet/finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed20+finetune0.005-learnable-indice0-heads4-shared1-routed0-topk0-mlpFalse-seed40.json @@ -0,0 +1,26 @@ +{ + "Naive": { + "Test Loss": [ + 0.8438, + 1.8672, + 0.9102 + ], + "Test Acc": [ + 72.2656, + 46.4844, + 68.75 + ] + }, + "permu_head_init_ortho_opt": { + "Test Loss": [ + 0.8281, + 1.0469, + 0.8828 + ], + "Test Acc": [ + 72.6562, + 64.4531, + 68.3594 + ] + } +} \ No newline at end of file diff --git a/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 b/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 new file mode 100644 index 0000000000000000000000000000000000000000..164355b5ad75af760a2c904293ff2f9e78564317 --- /dev/null +++ b/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 @@ -0,0 +1,115 @@ +{ + "train_loss_interp_naive": [ + 0.4399600926041603, + 0.4369426091015339, + 0.4334904816746712, + 0.4297081813216209, + 0.4256985554099083, + 0.4216733066737652, + 0.4178378953039646, + 0.4144207888841629, + 0.4115923511981964, + 0.40945807442069054, + 0.4080779683589935, + 0.40741123095154763, + 0.4074897059798241, + 0.40827907726168633, + 0.40972683936357496, + 0.41184418112039567, + 0.41459135577082634, + 0.4178251393139362, + 0.4214025369286537, + 0.42510213047266004, + 0.4287269660830498, + 0.43209666058421137, + 0.43507598310709, + 0.4376176556944847, + 0.4397105856239796 + ], + "test_loss_interp_naive": [ + 0.4399600926041603, + 0.4369426091015339, + 0.4334904816746712, + 0.4297081813216209, + 0.4256985554099083, + 0.4216733066737652, + 0.4178378953039646, + 0.4144207888841629, + 0.4115923511981964, + 0.40945807442069054, + 0.4080779683589935, + 0.40741123095154763, + 0.4074897059798241, + 0.40827907726168633, + 0.40972683936357496, + 0.41184418112039567, + 0.41459135577082634, + 0.4178251393139362, + 0.4214025369286537, + 0.42510213047266004, + 0.4287269660830498, + 0.43209666058421137, + 0.43507598310709, + 0.4376176556944847, + 0.4397105856239796 + ], + "train_acc_interp_naive": [ + 0.826, + 0.8262, + 0.8265, + 0.8267, + 0.8265, + 0.8259, + 0.8264, + 0.8254, + 0.8251, + 0.8245, + 0.8238, + 0.8232, + 0.8235, + 0.8235, + 0.8236, + 0.824, + 0.824, + 0.8254, + 0.8258, + 0.8259, + 0.8261, + 0.8279, + 0.8275, + 0.8285, + 0.8285 + ], + "test_acc_interp_naive": [ + 0.826, + 0.8262, + 0.8265, + 0.8267, + 0.8265, + 0.8259, + 0.8264, + 0.8254, + 0.8251, + 0.8245, + 0.8238, + 0.8232, + 0.8235, + 0.8235, + 0.8236, + 0.824, + 0.824, + 0.8254, + 0.8258, + 0.8259, + 0.8261, + 0.8279, + 0.8275, + 0.8285, + 0.8285 + ], + "train_loss_interp_clever_list": [], + "test_loss_interp_clever_list": [], + "train_acc_interp_clever_list": [], + "test_acc_interp_clever_list": [], + "baseline_train_loss": 0.4398353391140699 +} \ No newline at end of file diff --git a/results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20].json b/results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..d307f33b307c237d78224fbff39a8bf1b3e13e67 --- /dev/null +++ b/results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.4343, + 3.4369, + 3.4507, + 3.4792, + 3.5366, + 3.6814, + 4.1066, + 5.0711, + 6.321, + 7.304, + 7.8797, + 8.1689, + 8.2834, + 8.241, + 8.0033, + 7.5021, + 6.6494, + 5.4568, + 4.2838, + 3.7034, + 3.537, + 3.4807, + 3.4525, + 3.439, + 3.4362 + ], + "Test Loss": [ + 3.4343, + 3.4369, + 3.4507, + 3.4792, + 3.5366, + 3.6814, + 4.1066, + 5.0711, + 6.321, + 7.304, + 7.8797, + 8.1689, + 8.2834, + 8.241, + 8.0033, + 7.5021, + 6.6494, + 5.4568, + 4.2838, + 3.7034, + 3.537, + 3.4807, + 3.4525, + 3.439, + 3.4362 + ], + "Val PPL": [ + 31.0107, + 31.0893, + 31.5229, + 32.4348, + 34.3496, + 39.7011, + 60.7425, + 159.3499, + 556.1382, + 1486.2931, + 2643.0271, + 3529.3267, + 3957.7883, + 3793.4048, + 2990.7136, + 1811.8589, + 772.3456, + 234.3391, + 72.5145, + 40.5861, + 34.3647, + 32.4809, + 31.5785, + 31.1564, + 31.0672 + ], + "Test PPL": [ + 31.0107, + 31.0893, + 31.5229, + 32.4348, + 34.3496, + 39.7011, + 60.7425, + 159.3499, + 556.1382, + 1486.2931, + 2643.0271, + 3529.3267, + 3957.7883, + 3793.4048, + 2990.7136, + 1811.8589, + 772.3456, + 234.3391, + 72.5145, + 40.5861, + 34.3647, + 32.4809, + 31.5785, + 31.1564, + 31.0672 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.4343, + 3.4323, + 3.434, + 3.4385, + 3.4459, + 3.4572, + 3.4756, + 3.5079, + 3.5636, + 3.6435, + 3.732, + 3.801, + 3.8353, + 3.821, + 3.7634, + 3.6755, + 3.5852, + 3.5181, + 3.4797, + 3.4604, + 3.4485, + 3.4404, + 3.4359, + 3.4346, + 3.4364 + ], + "Test Loss": [ + 3.4343, + 3.4323, + 3.434, + 3.4385, + 3.4459, + 3.4572, + 3.4756, + 3.5079, + 3.5636, + 3.6435, + 3.732, + 3.801, + 3.8353, + 3.821, + 3.7634, + 3.6755, + 3.5852, + 3.5181, + 3.4797, + 3.4604, + 3.4485, + 3.4404, + 3.4359, + 3.4346, + 3.4364 + ], + "Val PPL": [ + 31.0107, + 30.9487, + 30.9991, + 31.1387, + 31.3701, + 31.7283, + 32.3161, + 33.3768, + 35.2911, + 38.2265, + 41.7642, + 44.7479, + 46.3084, + 45.6518, + 43.0938, + 39.4688, + 36.0612, + 33.7211, + 32.4506, + 31.8286, + 31.4518, + 31.1997, + 31.0607, + 31.0175, + 31.0763 + ], + "Test PPL": [ + 31.0107, + 30.9487, + 30.9991, + 31.1387, + 31.3701, + 31.7283, + 32.3161, + 33.3768, + 35.2911, + 38.2265, + 41.7642, + 44.7479, + 46.3084, + 45.6518, + 43.0938, + 39.4688, + 36.0612, + 33.7211, + 32.4506, + 31.8286, + 31.4518, + 31.1997, + 31.0607, + 31.0175, + 31.0763 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.4343, + 3.4328, + 3.4343, + 3.4386, + 3.446, + 3.4569, + 3.4727, + 3.4953, + 3.5317, + 3.5849, + 3.6527, + 3.7195, + 3.7649, + 3.7753, + 3.7456, + 3.6818, + 3.605, + 3.5363, + 3.4914, + 3.4663, + 3.4514, + 3.4417, + 3.4362, + 3.4344, + 3.4361 + ], + "Test Loss": [ + 3.4343, + 3.4328, + 3.4343, + 3.4386, + 3.446, + 3.4569, + 3.4727, + 3.4953, + 3.5317, + 3.5849, + 3.6527, + 3.7195, + 3.7649, + 3.7753, + 3.7456, + 3.6818, + 3.605, + 3.5363, + 3.4914, + 3.4663, + 3.4514, + 3.4417, + 3.4362, + 3.4344, + 3.4361 + ], + "Val PPL": [ + 31.0107, + 30.9617, + 31.0112, + 31.1424, + 31.3732, + 31.7171, + 32.2221, + 32.9613, + 34.1835, + 36.0492, + 38.5804, + 41.2447, + 43.1591, + 43.6092, + 42.3342, + 39.7177, + 36.7834, + 34.34, + 32.8323, + 32.019, + 31.5432, + 31.2413, + 31.0671, + 31.0127, + 31.0669 + ], + "Test PPL": [ + 31.0107, + 30.9617, + 31.0112, + 31.1424, + 31.3732, + 31.7171, + 32.2221, + 32.9613, + 34.1835, + 36.0492, + 38.5804, + 41.2447, + 43.1591, + 43.6092, + 42.3342, + 39.7177, + 36.7834, + 34.34, + 32.8323, + 32.019, + 31.5432, + 31.2413, + 31.0671, + 31.0127, + 31.0669 + ] + } +} \ No newline at end of file diff --git a/results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40].json b/results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..d10afa1802055699d9e4bc1a8f20a7a5ccef7d5d --- /dev/null +++ b/results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.4343, + 3.4371, + 3.45, + 3.4766, + 3.5351, + 3.7203, + 4.2566, + 5.3594, + 6.6144, + 7.4454, + 8.0407, + 8.3874, + 8.5146, + 8.5008, + 8.3437, + 7.9702, + 7.1617, + 5.9665, + 4.5642, + 3.7487, + 3.5321, + 3.4738, + 3.4468, + 3.4341, + 3.4321 + ], + "Test Loss": [ + 3.4343, + 3.4371, + 3.45, + 3.4766, + 3.5351, + 3.7203, + 4.2566, + 5.3594, + 6.6144, + 7.4454, + 8.0407, + 8.3874, + 8.5146, + 8.5008, + 8.3437, + 7.9702, + 7.1617, + 5.9665, + 4.5642, + 3.7487, + 3.5321, + 3.4738, + 3.4468, + 3.4341, + 3.4321 + ], + "Val PPL": [ + 31.0107, + 31.097, + 31.4997, + 32.351, + 34.2996, + 41.2759, + 70.5673, + 212.599, + 745.7325, + 1712.0389, + 3104.8508, + 4391.5137, + 4986.9028, + 4918.6699, + 4203.6479, + 2893.4844, + 1289.115, + 390.1571, + 95.9841, + 42.4667, + 34.1961, + 32.2585, + 31.4008, + 31.0046, + 30.9408 + ], + "Test PPL": [ + 31.0107, + 31.097, + 31.4997, + 32.351, + 34.2996, + 41.2759, + 70.5673, + 212.599, + 745.7325, + 1712.0389, + 3104.8508, + 4391.5137, + 4986.9028, + 4918.6699, + 4203.6479, + 2893.4844, + 1289.115, + 390.1571, + 95.9841, + 42.4667, + 34.1961, + 32.2585, + 31.4008, + 31.0046, + 30.9408 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.4343, + 3.4323, + 3.4338, + 3.4385, + 3.4469, + 3.46, + 3.4859, + 3.5529, + 3.6882, + 3.8601, + 4.0191, + 4.128, + 4.1693, + 4.1361, + 4.0365, + 3.8851, + 3.7135, + 3.5719, + 3.4923, + 3.4608, + 3.4464, + 3.4377, + 3.4324, + 3.4307, + 3.4322 + ], + "Test Loss": [ + 3.4343, + 3.4323, + 3.4338, + 3.4385, + 3.4469, + 3.46, + 3.4859, + 3.5529, + 3.6882, + 3.8601, + 4.0191, + 4.128, + 4.1693, + 4.1361, + 4.0365, + 3.8851, + 3.7135, + 3.5719, + 3.4923, + 3.4608, + 3.4464, + 3.4377, + 3.4324, + 3.4307, + 3.4322 + ], + "Val PPL": [ + 31.0107, + 30.9483, + 30.9928, + 31.1387, + 31.4021, + 31.8173, + 32.6526, + 34.9147, + 39.9726, + 47.4678, + 55.6533, + 62.0539, + 64.668, + 62.5562, + 56.6298, + 48.6699, + 40.9972, + 35.5838, + 32.861, + 31.8434, + 31.3878, + 31.1145, + 30.9519, + 30.899, + 30.9451 + ], + "Test PPL": [ + 31.0107, + 30.9483, + 30.9928, + 31.1387, + 31.4021, + 31.8173, + 32.6526, + 34.9147, + 39.9726, + 47.4678, + 55.6533, + 62.0539, + 64.668, + 62.5562, + 56.6298, + 48.6699, + 40.9972, + 35.5838, + 32.861, + 31.8434, + 31.3878, + 31.1145, + 30.9519, + 30.899, + 30.9451 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.4343, + 3.4327, + 3.4338, + 3.4381, + 3.4458, + 3.4572, + 3.4754, + 3.5136, + 3.5996, + 3.7351, + 3.8824, + 4.0031, + 4.0746, + 4.0784, + 4.017, + 3.8998, + 3.7493, + 3.6066, + 3.5101, + 3.4668, + 3.4486, + 3.4385, + 3.4327, + 3.4305, + 3.432 + ], + "Test Loss": [ + 3.4343, + 3.4327, + 3.4338, + 3.4381, + 3.4458, + 3.4572, + 3.4754, + 3.5136, + 3.5996, + 3.7351, + 3.8824, + 4.0031, + 4.0746, + 4.0784, + 4.017, + 3.8998, + 3.7493, + 3.6066, + 3.5101, + 3.4668, + 3.4486, + 3.4385, + 3.4327, + 3.4305, + 3.432 + ], + "Val PPL": [ + 31.0107, + 30.9591, + 30.995, + 31.1264, + 31.3668, + 31.7286, + 32.3099, + 33.5678, + 36.5845, + 41.8929, + 48.542, + 54.7686, + 58.8285, + 59.048, + 55.534, + 49.3923, + 42.4912, + 36.8404, + 33.4508, + 32.0326, + 31.4559, + 31.139, + 30.9607, + 30.891, + 30.9394 + ], + "Test PPL": [ + 31.0107, + 30.9591, + 30.995, + 31.1264, + 31.3668, + 31.7286, + 32.3099, + 33.5678, + 36.5845, + 41.8929, + 48.542, + 54.7686, + 58.8285, + 59.048, + 55.534, + 49.3923, + 42.4912, + 36.8404, + 33.4508, + 32.0326, + 31.4559, + 31.139, + 30.9607, + 30.891, + 30.9394 + ] + } +} \ No newline at end of file diff --git a/results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40].json b/results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..a0f4f03c57e57872c4291fa7d33057adc015ece6 --- /dev/null +++ b/results/lm1b/[finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.4362, + 3.4369, + 3.4475, + 3.4704, + 3.5348, + 3.9355, + 5.0334, + 6.1686, + 6.9804, + 7.522, + 7.8777, + 8.0989, + 8.1876, + 8.1137, + 7.893, + 7.5261, + 6.9881, + 6.1597, + 4.9975, + 3.9184, + 3.5322, + 3.467, + 3.4438, + 3.4332, + 3.4321 + ], + "Test Loss": [ + 3.4362, + 3.4369, + 3.4475, + 3.4704, + 3.5348, + 3.9355, + 5.0334, + 6.1686, + 6.9804, + 7.522, + 7.8777, + 8.0989, + 8.1876, + 8.1137, + 7.893, + 7.5261, + 6.9881, + 6.1597, + 4.9975, + 3.9184, + 3.5322, + 3.467, + 3.4438, + 3.4332, + 3.4321 + ], + "Val PPL": [ + 31.0672, + 31.0918, + 31.4215, + 32.1511, + 34.289, + 51.186, + 153.4466, + 477.5383, + 1075.3116, + 1848.2316, + 2637.9165, + 3290.717, + 3595.9927, + 3339.8457, + 2678.4255, + 1855.8431, + 1083.6989, + 473.2734, + 148.0436, + 50.3214, + 34.1977, + 32.041, + 31.3051, + 30.9754, + 30.9408 + ], + "Test PPL": [ + 31.0672, + 31.0918, + 31.4215, + 32.1511, + 34.289, + 51.186, + 153.4466, + 477.5383, + 1075.3116, + 1848.2316, + 2637.9165, + 3290.717, + 3595.9927, + 3339.8457, + 2678.4255, + 1855.8431, + 1083.6989, + 473.2734, + 148.0436, + 50.3214, + 34.1977, + 32.041, + 31.3051, + 30.9754, + 30.9408 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.4362, + 3.4343, + 3.436, + 3.4411, + 3.4502, + 3.4647, + 3.4908, + 3.5512, + 3.6703, + 3.8323, + 3.9936, + 4.1142, + 4.1563, + 4.1195, + 4.0067, + 3.8386, + 3.6699, + 3.5437, + 3.4859, + 3.4608, + 3.4464, + 3.4372, + 3.4318, + 3.4304, + 3.4321 + ], + "Test Loss": [ + 3.4362, + 3.4343, + 3.436, + 3.4411, + 3.4502, + 3.4647, + 3.4908, + 3.5512, + 3.6703, + 3.8323, + 3.9936, + 4.1142, + 4.1563, + 4.1195, + 4.0067, + 3.8386, + 3.6699, + 3.5437, + 3.4859, + 3.4608, + 3.4464, + 3.4372, + 3.4318, + 3.4304, + 3.4321 + ], + "Val PPL": [ + 31.0672, + 31.0104, + 31.0615, + 31.2226, + 31.5075, + 31.9661, + 32.812, + 34.855, + 39.2655, + 46.1698, + 54.2498, + 61.2022, + 63.8372, + 61.53, + 54.9641, + 46.4608, + 39.2477, + 34.5964, + 32.6516, + 31.8419, + 31.3859, + 31.1006, + 30.9318, + 30.8892, + 30.9427 + ], + "Test PPL": [ + 31.0672, + 31.0104, + 31.0615, + 31.2226, + 31.5075, + 31.9661, + 32.812, + 34.855, + 39.2655, + 46.1698, + 54.2498, + 61.2022, + 63.8372, + 61.53, + 54.9641, + 46.4608, + 39.2477, + 34.5964, + 32.6516, + 31.8419, + 31.3859, + 31.1006, + 30.9318, + 30.8892, + 30.9427 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.4362, + 3.435, + 3.4368, + 3.4419, + 3.4513, + 3.465, + 3.4861, + 3.5249, + 3.6012, + 3.7212, + 3.8591, + 3.9839, + 4.0587, + 4.0683, + 4.004, + 3.8808, + 3.7296, + 3.5878, + 3.5052, + 3.4687, + 3.4503, + 3.4389, + 3.4324, + 3.4304, + 3.4321 + ], + "Test Loss": [ + 3.4362, + 3.435, + 3.4368, + 3.4419, + 3.4513, + 3.465, + 3.4861, + 3.5249, + 3.6012, + 3.7212, + 3.8591, + 3.9839, + 4.0587, + 4.0683, + 4.004, + 3.8808, + 3.7296, + 3.5878, + 3.5052, + 3.4687, + 3.4503, + 3.4389, + 3.4324, + 3.4304, + 3.4321 + ], + "Val PPL": [ + 31.0672, + 31.0313, + 31.0879, + 31.2475, + 31.5399, + 31.9753, + 32.6576, + 33.9502, + 36.6417, + 41.3149, + 47.4242, + 53.725, + 57.9018, + 58.4602, + 54.817, + 48.4646, + 41.664, + 36.1551, + 33.2882, + 32.0959, + 31.5089, + 31.1531, + 30.9501, + 30.8877, + 30.9427 + ], + "Test PPL": [ + 31.0672, + 31.0313, + 31.0879, + 31.2475, + 31.5399, + 31.9753, + 32.6576, + 33.9502, + 36.6417, + 41.3149, + 47.4242, + 53.725, + 57.9018, + 58.4602, + 54.817, + 48.4646, + 41.664, + 36.1551, + 33.2882, + 32.0959, + 31.5089, + 31.1531, + 30.9501, + 30.8877, + 30.9427 + ] + } +} \ No newline at end of file diff --git a/results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json b/results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..f5f983ba31a923660f5655aa39e47d1462c6d0cd --- /dev/null +++ b/results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.4337, + 3.4353, + 3.4466, + 3.4708, + 3.5199, + 3.669, + 4.3133, + 5.6757, + 6.9758, + 7.8841, + 8.6178, + 9.1205, + 9.2884, + 9.1643, + 8.6724, + 7.9537, + 7.0626, + 5.8359, + 4.4564, + 3.6993, + 3.5246, + 3.472, + 3.4467, + 3.4346, + 3.4329 + ], + "Test Loss": [ + 3.4337, + 3.4353, + 3.4466, + 3.4708, + 3.5199, + 3.669, + 4.3133, + 5.6757, + 6.9758, + 7.8841, + 8.6178, + 9.1205, + 9.2884, + 9.1643, + 8.6724, + 7.9537, + 7.0626, + 5.8359, + 4.4564, + 3.6993, + 3.5246, + 3.472, + 3.4467, + 3.4346, + 3.4329 + ], + "Val PPL": [ + 30.9916, + 31.0403, + 31.3924, + 32.1614, + 33.7804, + 39.2136, + 74.687, + 291.6979, + 1070.3739, + 2654.6611, + 5529.1602, + 9140.4277, + 10811.9707, + 9549.8701, + 5839.7012, + 2846.0454, + 1167.5079, + 342.3635, + 86.1755, + 40.42, + 33.9386, + 32.2021, + 31.396, + 31.0199, + 30.9671 + ], + "Test PPL": [ + 30.9916, + 31.0403, + 31.3924, + 32.1614, + 33.7804, + 39.2136, + 74.687, + 291.6979, + 1070.3739, + 2654.6611, + 5529.1602, + 9140.4277, + 10811.9707, + 9549.8701, + 5839.7012, + 2846.0454, + 1167.5079, + 342.3635, + 86.1755, + 40.42, + 33.9386, + 32.2021, + 31.396, + 31.0199, + 30.9671 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.4337, + 3.4317, + 3.4326, + 3.4372, + 3.445, + 3.4576, + 3.4803, + 3.5374, + 3.6569, + 3.8337, + 4.0093, + 4.1358, + 4.1896, + 4.1511, + 4.0408, + 3.8714, + 3.6912, + 3.5553, + 3.4863, + 3.4586, + 3.445, + 3.4367, + 3.4322, + 3.4309, + 3.4329 + ], + "Test Loss": [ + 3.4337, + 3.4317, + 3.4326, + 3.4372, + 3.445, + 3.4576, + 3.4803, + 3.5374, + 3.6569, + 3.8337, + 4.0093, + 4.1358, + 4.1896, + 4.1511, + 4.0408, + 3.8714, + 3.6912, + 3.5553, + 3.4863, + 3.4586, + 3.445, + 3.4367, + 3.4322, + 3.4309, + 3.4329 + ], + "Val PPL": [ + 30.9916, + 30.9282, + 30.9581, + 31.0987, + 31.343, + 31.7416, + 32.4694, + 34.379, + 38.7413, + 46.232, + 55.1064, + 62.541, + 65.9992, + 63.5027, + 56.8717, + 48.0112, + 40.0937, + 34.9971, + 32.665, + 31.7714, + 31.3446, + 31.0838, + 30.9442, + 30.905, + 30.9659 + ], + "Test PPL": [ + 30.9916, + 30.9282, + 30.9581, + 31.0987, + 31.343, + 31.7416, + 32.4694, + 34.379, + 38.7413, + 46.232, + 55.1064, + 62.541, + 65.9992, + 63.5027, + 56.8717, + 48.0112, + 40.0937, + 34.9971, + 32.665, + 31.7714, + 31.3446, + 31.0838, + 30.9442, + 30.905, + 30.9659 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.4337, + 3.4322, + 3.4332, + 3.4372, + 3.4443, + 3.4545, + 3.4696, + 3.4929, + 3.5365, + 3.6138, + 3.7173, + 3.8247, + 3.9018, + 3.9182, + 3.8735, + 3.7759, + 3.6571, + 3.5536, + 3.4918, + 3.4632, + 3.4477, + 3.438, + 3.4325, + 3.4309, + 3.4329 + ], + "Test Loss": [ + 3.4337, + 3.4322, + 3.4332, + 3.4372, + 3.4443, + 3.4545, + 3.4696, + 3.4929, + 3.5365, + 3.6138, + 3.7173, + 3.8247, + 3.9018, + 3.9182, + 3.8735, + 3.7759, + 3.6571, + 3.5536, + 3.4918, + 3.4632, + 3.4477, + 3.438, + 3.4325, + 3.4309, + 3.4329 + ], + "Val PPL": [ + 30.9916, + 30.9436, + 30.9771, + 31.0991, + 31.3211, + 31.641, + 32.1229, + 32.8798, + 34.3475, + 37.1052, + 41.1544, + 45.8198, + 49.4928, + 50.3117, + 48.1117, + 43.6385, + 38.7483, + 34.9384, + 32.8465, + 31.9197, + 31.4268, + 31.1254, + 30.9554, + 30.9041, + 30.9669 + ], + "Test PPL": [ + 30.9916, + 30.9436, + 30.9771, + 31.0991, + 31.3211, + 31.641, + 32.1229, + 32.8798, + 34.3475, + 37.1052, + 41.1544, + 45.8198, + 49.4928, + 50.3117, + 48.1117, + 43.6385, + 38.7483, + 34.9384, + 32.8465, + 31.9197, + 31.4268, + 31.1254, + 30.9554, + 30.9041, + 30.9669 + ] + } +} \ No newline at end of file diff --git a/results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json b/results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..7a28c0b801f637a495e9f102ea13491dd600632c --- /dev/null +++ b/results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.4337, + 3.4353, + 3.4466, + 3.4715, + 3.5267, + 3.7117, + 4.3842, + 5.7822, + 7.3178, + 8.3845, + 8.8474, + 9.0178, + 9.0622, + 9.016, + 8.855, + 8.4706, + 7.5606, + 6.1351, + 4.7093, + 3.8526, + 3.551, + 3.4748, + 3.447, + 3.4346, + 3.4325 + ], + "Test Loss": [ + 3.4337, + 3.4353, + 3.4466, + 3.4715, + 3.5267, + 3.7117, + 4.3842, + 5.7822, + 7.3178, + 8.3845, + 8.8474, + 9.0178, + 9.0622, + 9.016, + 8.855, + 8.4706, + 7.5606, + 6.1351, + 4.7093, + 3.8526, + 3.551, + 3.4748, + 3.447, + 3.4346, + 3.4325 + ], + "Val PPL": [ + 30.9916, + 31.041, + 31.3923, + 32.1842, + 34.0125, + 40.9254, + 80.1733, + 324.4583, + 1506.9191, + 4378.8672, + 6956.1113, + 8248.5186, + 8622.875, + 8233.6797, + 7009.2979, + 4772.1782, + 1921.002, + 461.7938, + 110.9794, + 47.1157, + 34.8466, + 32.2923, + 31.4065, + 31.0188, + 30.9542 + ], + "Test PPL": [ + 30.9916, + 31.041, + 31.3923, + 32.1842, + 34.0125, + 40.9254, + 80.1733, + 324.4583, + 1506.9191, + 4378.8672, + 6956.1113, + 8248.5186, + 8622.875, + 8233.6797, + 7009.2979, + 4772.1782, + 1921.002, + 461.7938, + 110.9794, + 47.1157, + 34.8466, + 32.2923, + 31.4065, + 31.0188, + 30.9542 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.4337, + 3.4319, + 3.4339, + 3.4393, + 3.4484, + 3.4628, + 3.4881, + 3.5432, + 3.6522, + 3.8065, + 3.9665, + 4.0858, + 4.1269, + 4.0824, + 3.9639, + 3.8072, + 3.6545, + 3.5419, + 3.4855, + 3.4615, + 3.4475, + 3.4381, + 3.4326, + 3.431, + 3.4325 + ], + "Test Loss": [ + 3.4337, + 3.4319, + 3.4339, + 3.4393, + 3.4484, + 3.4628, + 3.4881, + 3.5432, + 3.6522, + 3.8065, + 3.9665, + 4.0858, + 4.1269, + 4.0824, + 3.9639, + 3.8072, + 3.6545, + 3.5419, + 3.4855, + 3.4615, + 3.4475, + 3.4381, + 3.4326, + 3.431, + 3.4325 + ], + "Val PPL": [ + 30.9916, + 30.9367, + 30.996, + 31.1637, + 31.4504, + 31.9056, + 32.7235, + 34.5762, + 38.5588, + 44.9915, + 52.8, + 59.4909, + 61.9886, + 59.2893, + 52.6627, + 45.0232, + 38.6496, + 34.5312, + 32.6377, + 31.8662, + 31.4221, + 31.1276, + 30.9585, + 30.9064, + 30.955 + ], + "Test PPL": [ + 30.9916, + 30.9367, + 30.996, + 31.1637, + 31.4504, + 31.9056, + 32.7235, + 34.5762, + 38.5588, + 44.9915, + 52.8, + 59.4909, + 61.9886, + 59.2893, + 52.6627, + 45.0232, + 38.6496, + 34.5312, + 32.6377, + 31.8662, + 31.4221, + 31.1276, + 30.9585, + 30.9064, + 30.955 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.4337, + 3.4323, + 3.434, + 3.4389, + 3.4472, + 3.4598, + 3.4791, + 3.5134, + 3.5781, + 3.6796, + 3.8066, + 3.9289, + 4.0066, + 4.0121, + 3.9456, + 3.8255, + 3.6876, + 3.5692, + 3.4983, + 3.4671, + 3.4499, + 3.4392, + 3.4328, + 3.4308, + 3.4326 + ], + "Test Loss": [ + 3.4337, + 3.4323, + 3.434, + 3.4389, + 3.4472, + 3.4598, + 3.4791, + 3.5134, + 3.5781, + 3.6796, + 3.8066, + 3.9289, + 4.0066, + 4.0121, + 3.9456, + 3.8255, + 3.6876, + 3.5692, + 3.4983, + 3.4671, + 3.4499, + 3.4392, + 3.4328, + 3.4308, + 3.4326 + ], + "Val PPL": [ + 30.9916, + 30.9464, + 30.9989, + 31.1516, + 31.4119, + 31.8121, + 32.4292, + 33.5638, + 35.8038, + 39.6322, + 44.9994, + 50.8485, + 54.9582, + 55.2654, + 51.7058, + 45.8563, + 39.9476, + 35.4895, + 33.0585, + 32.0428, + 31.4984, + 31.1605, + 30.9641, + 30.9013, + 30.958 + ], + "Test PPL": [ + 30.9916, + 30.9464, + 30.9989, + 31.1516, + 31.4119, + 31.8121, + 32.4292, + 33.5638, + 35.8038, + 39.6322, + 44.9994, + 50.8485, + 54.9582, + 55.2654, + 51.7058, + 45.8563, + 39.9476, + 35.4895, + 33.0585, + 32.0428, + 31.4984, + 31.1605, + 30.9641, + 30.9013, + 30.958 + ] + } +} \ No newline at end of file diff --git a/results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json b/results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..fcfb39b9bea3bf843c6591d3d452eed8490b9569 --- /dev/null +++ b/results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.4329, + 3.4336, + 3.4448, + 3.4711, + 3.5614, + 4.1672, + 5.5089, + 6.8289, + 7.7999, + 8.4597, + 8.8539, + 9.0535, + 9.104, + 9.0117, + 8.764, + 8.324, + 7.6937, + 6.697, + 5.2721, + 3.9815, + 3.5402, + 3.4714, + 3.4456, + 3.4341, + 3.4325 + ], + "Test Loss": [ + 3.4329, + 3.4336, + 3.4448, + 3.4711, + 3.5614, + 4.1672, + 5.5089, + 6.8289, + 7.7999, + 8.4597, + 8.8539, + 9.0535, + 9.104, + 9.0117, + 8.764, + 8.324, + 7.6937, + 6.697, + 5.2721, + 3.9815, + 3.5402, + 3.4714, + 3.4456, + 3.4341, + 3.4325 + ], + "Val PPL": [ + 30.9671, + 30.9887, + 31.3356, + 32.1728, + 35.2109, + 64.534, + 246.8717, + 924.1573, + 2440.4697, + 4720.811, + 7001.4551, + 8548.749, + 8991.2891, + 8198.7881, + 6399.4937, + 4121.6978, + 2194.5425, + 809.9375, + 194.8281, + 53.5994, + 34.475, + 32.1813, + 31.3609, + 31.0028, + 30.9542 + ], + "Test PPL": [ + 30.9671, + 30.9887, + 31.3356, + 32.1728, + 35.2109, + 64.534, + 246.8717, + 924.1573, + 2440.4697, + 4720.811, + 7001.4551, + 8548.749, + 8991.2891, + 8198.7881, + 6399.4937, + 4121.6978, + 2194.5425, + 809.9375, + 194.8281, + 53.5994, + 34.475, + 32.1813, + 31.3609, + 31.0028, + 30.9542 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.4329, + 3.4311, + 3.4329, + 3.4383, + 3.4476, + 3.4638, + 3.5039, + 3.6106, + 3.7918, + 4.0092, + 4.2008, + 4.3263, + 4.3639, + 4.3041, + 4.1528, + 3.9473, + 3.7358, + 3.5734, + 3.4911, + 3.4602, + 3.4459, + 3.4374, + 3.4322, + 3.4308, + 3.4324 + ], + "Test Loss": [ + 3.4329, + 3.4311, + 3.4329, + 3.4383, + 3.4476, + 3.4638, + 3.5039, + 3.6106, + 3.7918, + 4.0092, + 4.2008, + 4.3263, + 4.3639, + 4.3041, + 4.1528, + 3.9473, + 3.7358, + 3.5734, + 3.4911, + 3.4602, + 3.4459, + 3.4374, + 3.4322, + 3.4308, + 3.4324 + ], + "Val PPL": [ + 30.9671, + 30.9113, + 30.9657, + 31.134, + 31.4236, + 31.9387, + 33.2447, + 36.9869, + 44.3378, + 55.1052, + 66.7398, + 75.6654, + 78.5596, + 74.0042, + 63.6146, + 51.7936, + 41.9208, + 35.6391, + 32.8231, + 31.8246, + 31.3723, + 31.106, + 30.9448, + 30.9013, + 30.9515 + ], + "Test PPL": [ + 30.9671, + 30.9113, + 30.9657, + 31.134, + 31.4236, + 31.9387, + 33.2447, + 36.9869, + 44.3378, + 55.1052, + 66.7398, + 75.6654, + 78.5596, + 74.0042, + 63.6146, + 51.7936, + 41.9208, + 35.6391, + 32.8231, + 31.8246, + 31.3723, + 31.106, + 30.9448, + 30.9013, + 30.9515 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.4329, + 3.4315, + 3.4332, + 3.4381, + 3.4466, + 3.4596, + 3.4851, + 3.5491, + 3.6652, + 3.8134, + 3.9633, + 4.0869, + 4.158, + 4.1595, + 4.0838, + 3.9505, + 3.7903, + 3.6329, + 3.5182, + 3.4674, + 3.4488, + 3.4384, + 3.4327, + 3.4308, + 3.4323 + ], + "Test Loss": [ + 3.4329, + 3.4315, + 3.4332, + 3.4381, + 3.4466, + 3.4596, + 3.4851, + 3.5491, + 3.6652, + 3.8134, + 3.9633, + 4.0869, + 4.158, + 4.1595, + 4.0838, + 3.9505, + 3.7903, + 3.6329, + 3.5182, + 3.4674, + 3.4488, + 3.4384, + 3.4327, + 3.4308, + 3.4323 + ], + "Val PPL": [ + 30.9671, + 30.9222, + 30.9756, + 31.127, + 31.3938, + 31.8039, + 32.6264, + 34.7826, + 39.065, + 45.3057, + 52.6293, + 59.5547, + 63.9459, + 64.0391, + 59.3685, + 51.9631, + 44.2693, + 37.8221, + 33.722, + 32.0545, + 31.4635, + 31.1386, + 30.9599, + 30.901, + 30.9493 + ], + "Test PPL": [ + 30.9671, + 30.9222, + 30.9756, + 31.127, + 31.3938, + 31.8039, + 32.6264, + 34.7826, + 39.065, + 45.3057, + 52.6293, + 59.5547, + 63.9459, + 64.0391, + 59.3685, + 51.9631, + 44.2693, + 37.8221, + 33.722, + 32.0545, + 31.4635, + 31.1386, + 30.9599, + 30.901, + 30.9493 + ] + } +} \ No newline at end of file diff --git a/results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json b/results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..723656a1827cdd7599948235f19ece4cf87730e2 --- /dev/null +++ b/results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.4389, + 3.4416, + 3.4555, + 3.4852, + 3.5584, + 3.8725, + 4.813, + 6.0781, + 7.0912, + 7.6963, + 8.1042, + 8.3769, + 8.4825, + 8.4201, + 8.1905, + 7.8067, + 7.202, + 6.2474, + 4.9833, + 3.9547, + 3.5717, + 3.4889, + 3.4561, + 3.4405, + 3.4373 + ], + "Test Loss": [ + 3.4389, + 3.4416, + 3.4555, + 3.4852, + 3.5584, + 3.8725, + 4.813, + 6.0781, + 7.0912, + 7.6963, + 8.1042, + 8.3769, + 8.4825, + 8.4201, + 8.1905, + 7.8067, + 7.202, + 6.2474, + 4.9833, + 3.9547, + 3.5717, + 3.4889, + 3.4561, + 3.4405, + 3.4373 + ], + "Val PPL": [ + 31.1514, + 31.2372, + 31.6731, + 32.6303, + 35.1053, + 48.0626, + 123.1003, + 436.2155, + 1201.3105, + 2200.272, + 3308.4016, + 4345.6982, + 4829.5459, + 4537.583, + 3606.5845, + 2457.0693, + 1342.1694, + 516.6775, + 145.9558, + 52.1775, + 35.5784, + 32.7502, + 31.6944, + 31.2032, + 31.1015 + ], + "Test PPL": [ + 31.1514, + 31.2372, + 31.6731, + 32.6303, + 35.1053, + 48.0626, + 123.1003, + 436.2155, + 1201.3105, + 2200.272, + 3308.4016, + 4345.6982, + 4829.5459, + 4537.583, + 3606.5845, + 2457.0693, + 1342.1694, + 516.6775, + 145.9558, + 52.1775, + 35.5784, + 32.7502, + 31.6944, + 31.2032, + 31.1015 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.4389, + 3.4372, + 3.4379, + 3.4406, + 3.4454, + 3.4529, + 3.4643, + 3.4886, + 3.5338, + 3.5952, + 3.658, + 3.7055, + 3.7263, + 3.7153, + 3.675, + 3.6101, + 3.5404, + 3.4887, + 3.4638, + 3.4519, + 3.4444, + 3.4393, + 3.4362, + 3.4355, + 3.4373 + ], + "Test Loss": [ + 3.4389, + 3.4372, + 3.4379, + 3.4406, + 3.4454, + 3.4529, + 3.4643, + 3.4886, + 3.5338, + 3.5952, + 3.658, + 3.7055, + 3.7263, + 3.7153, + 3.675, + 3.6101, + 3.5404, + 3.4887, + 3.4638, + 3.4519, + 3.4444, + 3.4393, + 3.4362, + 3.4355, + 3.4373 + ], + "Val PPL": [ + 31.1514, + 31.1, + 31.1214, + 31.2061, + 31.3556, + 31.593, + 31.9544, + 32.7385, + 34.2541, + 36.4234, + 38.7828, + 40.6722, + 41.5246, + 41.0699, + 39.4503, + 36.9699, + 34.4794, + 32.7433, + 31.9371, + 31.5612, + 31.3238, + 31.1655, + 31.0683, + 31.0472, + 31.1019 + ], + "Test PPL": [ + 31.1514, + 31.1, + 31.1214, + 31.2061, + 31.3556, + 31.593, + 31.9544, + 32.7385, + 34.2541, + 36.4234, + 38.7828, + 40.6722, + 41.5246, + 41.0699, + 39.4503, + 36.9699, + 34.4794, + 32.7433, + 31.9371, + 31.5612, + 31.3238, + 31.1655, + 31.0683, + 31.0472, + 31.1019 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.4389, + 3.4376, + 3.4385, + 3.4414, + 3.4457, + 3.4523, + 3.4613, + 3.4744, + 3.4987, + 3.5378, + 3.5853, + 3.6301, + 3.6615, + 3.6679, + 3.6494, + 3.6052, + 3.5482, + 3.4985, + 3.4698, + 3.4554, + 3.4464, + 3.4401, + 3.4365, + 3.4354, + 3.4371 + ], + "Test Loss": [ + 3.4389, + 3.4376, + 3.4385, + 3.4414, + 3.4457, + 3.4523, + 3.4613, + 3.4744, + 3.4987, + 3.5378, + 3.5853, + 3.6301, + 3.6615, + 3.6679, + 3.6494, + 3.6052, + 3.5482, + 3.4985, + 3.4698, + 3.4554, + 3.4464, + 3.4401, + 3.4365, + 3.4354, + 3.4371 + ], + "Val PPL": [ + 31.1514, + 31.1131, + 31.1408, + 31.2299, + 31.3665, + 31.5724, + 31.858, + 32.2777, + 33.0718, + 34.3896, + 36.0627, + 37.7179, + 38.9212, + 39.1691, + 38.4534, + 36.7903, + 34.7501, + 33.067, + 32.1316, + 31.6708, + 31.388, + 31.1914, + 31.0777, + 31.0435, + 31.0978 + ], + "Test PPL": [ + 31.1514, + 31.1131, + 31.1408, + 31.2299, + 31.3665, + 31.5724, + 31.858, + 32.2777, + 33.0718, + 34.3896, + 36.0627, + 37.7179, + 38.9212, + 39.1691, + 38.4534, + 36.7903, + 34.7501, + 33.067, + 32.1316, + 31.6708, + 31.388, + 31.1914, + 31.0777, + 31.0435, + 31.0978 + ] + } +} \ No newline at end of file diff --git a/results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json b/results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..5acb5cdce910ef8fc0d1f663d1c72a3368bf919f --- /dev/null +++ b/results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.4389, + 3.4413, + 3.4545, + 3.4823, + 3.5445, + 3.7446, + 4.3328, + 5.422, + 6.5959, + 7.6509, + 8.2644, + 8.5024, + 8.563, + 8.4923, + 8.2343, + 7.6594, + 6.559, + 5.323, + 4.2856, + 3.7238, + 3.5405, + 3.4817, + 3.4537, + 3.4401, + 3.4373 + ], + "Test Loss": [ + 3.4389, + 3.4413, + 3.4545, + 3.4823, + 3.5445, + 3.7446, + 4.3328, + 5.422, + 6.5959, + 7.6509, + 8.2644, + 8.5024, + 8.563, + 8.4923, + 8.2343, + 7.6594, + 6.559, + 5.323, + 4.2856, + 3.7238, + 3.5405, + 3.4817, + 3.4537, + 3.4401, + 3.4373 + ], + "Val PPL": [ + 31.1514, + 31.2275, + 31.6413, + 32.534, + 34.6236, + 42.2927, + 76.156, + 226.3265, + 732.0634, + 2102.6301, + 3883.1072, + 4926.5522, + 5234.479, + 4877.2563, + 3768.0361, + 2120.5222, + 705.557, + 205.0042, + 72.6495, + 41.4235, + 34.484, + 32.5159, + 31.6172, + 31.1908, + 31.1032 + ], + "Test PPL": [ + 31.1514, + 31.2275, + 31.6413, + 32.534, + 34.6236, + 42.2927, + 76.156, + 226.3265, + 732.0634, + 2102.6301, + 3883.1072, + 4926.5522, + 5234.479, + 4877.2563, + 3768.0361, + 2120.5222, + 705.557, + 205.0042, + 72.6495, + 41.4235, + 34.484, + 32.5159, + 31.6172, + 31.1908, + 31.1032 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.4389, + 3.4371, + 3.438, + 3.4415, + 3.4478, + 3.4572, + 3.4737, + 3.5071, + 3.567, + 3.6447, + 3.7246, + 3.7903, + 3.8249, + 3.812, + 3.7564, + 3.6709, + 3.5818, + 3.5138, + 3.476, + 3.4582, + 3.4476, + 3.4408, + 3.4367, + 3.4355, + 3.4373 + ], + "Test Loss": [ + 3.4389, + 3.4371, + 3.438, + 3.4415, + 3.4478, + 3.4572, + 3.4737, + 3.5071, + 3.567, + 3.6447, + 3.7246, + 3.7903, + 3.8249, + 3.812, + 3.7564, + 3.6709, + 3.5818, + 3.5138, + 3.476, + 3.4582, + 3.4476, + 3.4408, + 3.4367, + 3.4355, + 3.4373 + ], + "Val PPL": [ + 31.1514, + 31.0968, + 31.1247, + 31.2337, + 31.4298, + 31.7295, + 32.2564, + 33.3513, + 35.4087, + 38.2728, + 41.4556, + 44.2699, + 45.8264, + 45.2414, + 42.7953, + 39.2878, + 35.9364, + 33.5767, + 32.3301, + 31.7584, + 31.4238, + 31.2133, + 31.0841, + 31.0475, + 31.1013 + ], + "Test PPL": [ + 31.1514, + 31.0968, + 31.1247, + 31.2337, + 31.4298, + 31.7295, + 32.2564, + 33.3513, + 35.4087, + 38.2728, + 41.4556, + 44.2699, + 45.8264, + 45.2414, + 42.7953, + 39.2878, + 35.9364, + 33.5767, + 32.3301, + 31.7584, + 31.4238, + 31.2133, + 31.0841, + 31.0475, + 31.1013 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.4389, + 3.4374, + 3.4385, + 3.442, + 3.4479, + 3.4565, + 3.4692, + 3.4887, + 3.5215, + 3.5728, + 3.6362, + 3.697, + 3.7398, + 3.749, + 3.7172, + 3.6545, + 3.5815, + 3.5198, + 3.4818, + 3.4616, + 3.4492, + 3.441, + 3.4367, + 3.4352, + 3.4372 + ], + "Test Loss": [ + 3.4389, + 3.4374, + 3.4385, + 3.442, + 3.4479, + 3.4565, + 3.4692, + 3.4887, + 3.5215, + 3.5728, + 3.6362, + 3.697, + 3.7398, + 3.749, + 3.7172, + 3.6545, + 3.5815, + 3.5198, + 3.4818, + 3.4616, + 3.4492, + 3.441, + 3.4367, + 3.4352, + 3.4372 + ], + "Val PPL": [ + 31.1514, + 31.1054, + 31.1402, + 31.2493, + 31.4351, + 31.7062, + 32.112, + 32.7437, + 33.8353, + 35.6157, + 37.9489, + 40.3243, + 42.0882, + 42.4765, + 41.1494, + 38.6468, + 35.926, + 33.778, + 32.518, + 31.8666, + 31.474, + 31.2172, + 31.0836, + 31.0368, + 31.0998 + ], + "Test PPL": [ + 31.1514, + 31.1054, + 31.1402, + 31.2493, + 31.4351, + 31.7062, + 32.112, + 32.7437, + 33.8353, + 35.6157, + 37.9489, + 40.3243, + 42.0882, + 42.4765, + 41.1494, + 38.6468, + 35.926, + 33.778, + 32.518, + 31.8666, + 31.474, + 31.2172, + 31.0836, + 31.0368, + 31.0998 + ] + } +} \ No newline at end of file diff --git a/results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json b/results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..b5c09fd79a02adc3a427d50a32e0c694bfeeb442 --- /dev/null +++ b/results/lm1b/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.4373, + 3.4394, + 3.4529, + 3.4844, + 3.5938, + 4.3001, + 5.8354, + 7.0974, + 7.9612, + 8.515, + 8.7596, + 8.8701, + 8.9061, + 8.8579, + 8.6976, + 8.3097, + 7.6131, + 6.7598, + 5.4153, + 4.0517, + 3.5624, + 3.4803, + 3.452, + 3.4392, + 3.4373 + ], + "Test Loss": [ + 3.4373, + 3.4394, + 3.4529, + 3.4844, + 3.5938, + 4.3001, + 5.8354, + 7.0974, + 7.9612, + 8.515, + 8.7596, + 8.8701, + 8.9061, + 8.8579, + 8.6976, + 8.3097, + 7.6131, + 6.7598, + 5.4153, + 4.0517, + 3.5624, + 3.4803, + 3.452, + 3.4392, + 3.4373 + ], + "Val PPL": [ + 31.1015, + 31.1679, + 31.5912, + 32.6014, + 36.3703, + 73.7105, + 342.1931, + 1208.8657, + 2867.384, + 4988.9717, + 6371.2568, + 7116.0879, + 7376.9639, + 7029.6226, + 5988.4448, + 4063.0515, + 2024.5107, + 862.4745, + 224.8237, + 57.4971, + 35.2474, + 32.4705, + 31.5628, + 31.1616, + 31.1032 + ], + "Test PPL": [ + 31.1015, + 31.1679, + 31.5912, + 32.6014, + 36.3703, + 73.7105, + 342.1931, + 1208.8657, + 2867.384, + 4988.9717, + 6371.2568, + 7116.0879, + 7376.9639, + 7029.6226, + 5988.4448, + 4063.0515, + 2024.5107, + 862.4745, + 224.8237, + 57.4971, + 35.2474, + 32.4705, + 31.5628, + 31.1616, + 31.1032 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.4373, + 3.4353, + 3.4361, + 3.4395, + 3.4458, + 3.4554, + 3.4704, + 3.4961, + 3.541, + 3.6017, + 3.6636, + 3.708, + 3.7217, + 3.7, + 3.6501, + 3.5888, + 3.5321, + 3.4935, + 3.4702, + 3.4561, + 3.4465, + 3.4402, + 3.4364, + 3.4355, + 3.4372 + ], + "Test Loss": [ + 3.4373, + 3.4353, + 3.4361, + 3.4395, + 3.4458, + 3.4554, + 3.4704, + 3.4961, + 3.541, + 3.6017, + 3.6636, + 3.708, + 3.7217, + 3.7, + 3.6501, + 3.5888, + 3.5321, + 3.4935, + 3.4702, + 3.4561, + 3.4465, + 3.4402, + 3.4364, + 3.4355, + 3.4372 + ], + "Val PPL": [ + 31.1015, + 31.0406, + 31.0658, + 31.1718, + 31.3684, + 31.6725, + 32.15, + 32.9854, + 34.5016, + 36.6596, + 39.0001, + 40.7713, + 41.3347, + 40.4464, + 38.4795, + 36.1922, + 34.1952, + 32.9021, + 32.1419, + 31.6925, + 31.3901, + 31.193, + 31.0752, + 31.0464, + 31.0984 + ], + "Test PPL": [ + 31.1015, + 31.0406, + 31.0658, + 31.1718, + 31.3684, + 31.6725, + 32.15, + 32.9854, + 34.5016, + 36.6596, + 39.0001, + 40.7713, + 41.3347, + 40.4464, + 38.4795, + 36.1922, + 34.1952, + 32.9021, + 32.1419, + 31.6925, + 31.3901, + 31.193, + 31.0752, + 31.0464, + 31.0984 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.4373, + 3.4355, + 3.4363, + 3.4394, + 3.4459, + 3.4555, + 3.4689, + 3.4887, + 3.5207, + 3.5702, + 3.6305, + 3.6865, + 3.7193, + 3.7182, + 3.683, + 3.6233, + 3.5603, + 3.5111, + 3.4804, + 3.4614, + 3.4493, + 3.4411, + 3.4364, + 3.4353, + 3.4371 + ], + "Test Loss": [ + 3.4373, + 3.4355, + 3.4363, + 3.4394, + 3.4459, + 3.4555, + 3.4689, + 3.4887, + 3.5207, + 3.5702, + 3.6305, + 3.6865, + 3.7193, + 3.7182, + 3.683, + 3.6233, + 3.5603, + 3.5111, + 3.4804, + 3.4614, + 3.4493, + 3.4411, + 3.4364, + 3.4353, + 3.4371 + ], + "Val PPL": [ + 31.1015, + 31.0476, + 31.0711, + 31.1698, + 31.3718, + 31.6727, + 32.101, + 32.7436, + 33.8088, + 35.5231, + 37.7309, + 39.9048, + 41.2373, + 41.1898, + 39.7662, + 37.4594, + 35.1734, + 33.4853, + 32.4733, + 31.8621, + 31.4784, + 31.2201, + 31.0755, + 31.0395, + 31.0978 + ], + "Test PPL": [ + 31.1015, + 31.0476, + 31.0711, + 31.1698, + 31.3718, + 31.6727, + 32.101, + 32.7436, + 33.8088, + 35.5231, + 37.7309, + 39.9048, + 41.2373, + 41.1898, + 39.7662, + 37.4594, + 35.1734, + 33.4853, + 32.4733, + 31.8621, + 31.4784, + 31.2201, + 31.0755, + 31.0395, + 31.0978 + ] + } +} \ No newline at end of file diff --git a/results/penn/[lr0.0002-topk2-shared0-routed2-seed0+lr0.0002-topk2-shared0-routed2-seed20].json b/results/penn/[lr0.0002-topk2-shared0-routed2-seed0+lr0.0002-topk2-shared0-routed2-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..74abbe8972ed175273cf1996c6c9f742df4d04ce --- /dev/null +++ b/results/penn/[lr0.0002-topk2-shared0-routed2-seed0+lr0.0002-topk2-shared0-routed2-seed20].json @@ -0,0 +1,335 @@ +{ + "train_loss_interp_naive": [ + 4.463920593261719, + 4.4646759033203125, + 4.467608451843262, + 4.472905158996582, + 4.480701446533203, + 4.490994453430176, + 4.50356388092041, + 4.517791748046875, + 4.532586097717285, + 4.546379089355469, + 4.557367324829102, + 4.563878536224365, + 4.5647969245910645, + 4.5598978996276855, + 4.549891948699951, + 4.536221027374268, + 4.520663261413574, + 4.504913806915283, + 4.490306854248047, + 4.47770881652832, + 4.467562675476074, + 4.460000038146973, + 4.4549479484558105, + 4.452245712280273, + 4.451686859130859 + ], + "test_loss_interp_naive": [ + 4.311151027679443, + 4.312248706817627, + 4.315517425537109, + 4.3211283683776855, + 4.329193115234375, + 4.339681148529053, + 4.352320671081543, + 4.366482734680176, + 4.381086349487305, + 4.394619941711426, + 4.405344486236572, + 4.4116435050964355, + 4.412441253662109, + 4.407537937164307, + 4.397638320922852, + 4.384158611297607, + 4.368839740753174, + 4.353326320648193, + 4.338901519775391, + 4.326379299163818, + 4.316190719604492, + 4.3084588050842285, + 4.303143501281738, + 4.300105094909668, + 4.299169063568115 + ], + "train_ppl_interp_naive": [ + 86.82725524902344, + 86.89286041259766, + 87.14805603027344, + 87.61087799072266, + 88.29658508300781, + 89.21011352539062, + 90.33851623535156, + 91.63302612304688, + 92.99874877929688, + 94.29037475585938, + 95.33216857910156, + 95.95492553710938, + 96.04308319091797, + 95.57371520996094, + 94.62217712402344, + 93.33740997314453, + 91.89653015136719, + 90.4605484008789, + 89.14878845214844, + 88.0327377319336, + 87.1440658569336, + 86.48751068115234, + 86.05167388916016, + 85.8194580078125, + 85.7715072631836 + ], + "test_ppl_interp_naive": [ + 74.5262222290039, + 74.60807800292969, + 74.85234069824219, + 75.27351379394531, + 75.8830337524414, + 76.68309020996094, + 77.65847778320312, + 78.76610565185547, + 79.92481231689453, + 81.01383972167969, + 81.8873519897461, + 82.40479278564453, + 82.47055053710938, + 82.0671615600586, + 81.25873565673828, + 80.17074584960938, + 78.95197296142578, + 77.73661041259766, + 76.62332153320312, + 75.66981506347656, + 74.90276336669922, + 74.32585144042969, + 73.93183898925781, + 73.7075424194336, + 73.63858032226562 + ], + "train_loss_interp_clever_list": [ + [ + 4.463920593261719, + 4.464303970336914, + 4.465911865234375, + 4.468739032745361, + 4.472715377807617, + 4.477677822113037, + 4.483362674713135, + 4.489373683929443, + 4.495208740234375, + 4.5002875328063965, + 4.50404167175293, + 4.5059943199157715, + 4.505858898162842, + 4.503597736358643, + 4.499417304992676, + 4.493732452392578, + 4.4870829582214355, + 4.480055332183838, + 4.473166465759277, + 4.466861248016357, + 4.46144962310791, + 4.4571404457092285, + 4.4540510177612305, + 4.452232360839844, + 4.451684951782227 + ], + [ + 4.463920593261719, + 4.4645304679870605, + 4.466420650482178, + 4.469551086425781, + 4.473820686340332, + 4.479038715362549, + 4.484922409057617, + 4.49106502532959, + 4.49696159362793, + 4.502043724060059, + 4.505760669708252, + 4.50766658782959, + 4.507490158081055, + 4.505192756652832, + 4.500972747802734, + 4.495230674743652, + 4.488480567932129, + 4.4812912940979, + 4.4741973876953125, + 4.467654228210449, + 4.461999893188477, + 4.4574713706970215, + 4.454208850860596, + 4.452276229858398, + 4.451685905456543 + ] + ], + "test_loss_interp_clever_list": [ + [ + 4.311151027679443, + 4.3117756843566895, + 4.31365442276001, + 4.3167724609375, + 4.3210554122924805, + 4.326333522796631, + 4.332325458526611, + 4.338619232177734, + 4.344696998596191, + 4.349958419799805, + 4.353827953338623, + 4.355823993682861, + 4.355664253234863, + 4.353318214416504, + 4.349003314971924, + 4.343141555786133, + 4.336293697357178, + 4.329051971435547, + 4.321949481964111, + 4.315425872802734, + 4.309799671173096, + 4.3052778244018555, + 4.301970958709717, + 4.299932956695557, + 4.299170017242432 + ], + [ + 4.311151027679443, + 4.3119587898254395, + 4.3140435218811035, + 4.317363262176514, + 4.321816444396973, + 4.327211856842041, + 4.333254337310791, + 4.339531421661377, + 4.34553861618042, + 4.350697040557861, + 4.354468822479248, + 4.356398105621338, + 4.356229782104492, + 4.353923797607422, + 4.349687576293945, + 4.343922138214111, + 4.337143421173096, + 4.3299150466918945, + 4.322765827178955, + 4.316141605377197, + 4.3103718757629395, + 4.305686950683594, + 4.3022236824035645, + 4.300046920776367, + 4.299169540405273 + ] + ], + "train_ppl_interp_clever_list": [ + [ + 86.82725524902344, + 86.86054992675781, + 87.00032043457031, + 87.24663543701172, + 87.59425354003906, + 88.03001403808594, + 88.5318832397461, + 89.06564331054688, + 89.58686828613281, + 90.04302215576172, + 90.3816909790039, + 90.55834197998047, + 90.54607391357422, + 90.34156799316406, + 89.9646987915039, + 89.4547119140625, + 88.86184692382812, + 88.23955535888672, + 87.63377380371094, + 87.08296203613281, + 86.61296844482422, + 86.24054718017578, + 85.97452545166016, + 85.81830596923828, + 85.7713394165039 + ], + [ + 86.82725524902344, + 86.8802261352539, + 87.04460144042969, + 87.31752014160156, + 87.69112396240234, + 88.14989471435547, + 88.67007446289062, + 89.2164077758789, + 89.74403381347656, + 90.20128631591797, + 90.53718566894531, + 90.70989990234375, + 90.69390106201172, + 90.48578643798828, + 90.104736328125, + 89.58883666992188, + 88.98612976074219, + 88.34867858886719, + 87.72415924072266, + 87.15204620361328, + 86.66064453125, + 86.26909637451172, + 85.98809814453125, + 85.82207489013672, + 85.77142333984375 + ] + ], + "test_ppl_interp_clever_list": [ + [ + 74.5262222290039, + 74.57279205322266, + 74.71302032470703, + 74.94634246826172, + 75.26802062988281, + 75.66635131835938, + 76.12110137939453, + 76.60169982910156, + 77.06867980957031, + 77.4752426147461, + 77.7756118774414, + 77.93101501464844, + 77.91856384277344, + 77.73597717285156, + 77.40127563476562, + 76.94889831542969, + 76.42376708984375, + 75.87232208251953, + 75.3353500366211, + 74.84548950195312, + 74.42557525634766, + 74.08979797363281, + 73.84519958496094, + 73.69485473632812, + 73.63865661621094 + ], + [ + 74.5262222290039, + 74.5864486694336, + 74.74210357666016, + 74.99063873291016, + 75.32533264160156, + 75.73284149169922, + 76.19184112548828, + 76.6716079711914, + 77.1335678100586, + 77.53248596191406, + 77.82547760009766, + 77.97576904296875, + 77.962646484375, + 77.78306579589844, + 77.45425415039062, + 77.00898742675781, + 76.48872375488281, + 75.93783569335938, + 75.3968734741211, + 74.89907836914062, + 74.46817779541016, + 74.1201171875, + 73.86386108398438, + 73.70325469970703, + 73.63861846923828 + ] + ], + "baseline_train_loss": 4.457803726196289 +} \ No newline at end of file diff --git a/results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json b/results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..ee9db942998e503aa99117330903730300352751 --- /dev/null +++ b/results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.7779, + 1.7799, + 1.7935, + 1.8252, + 1.8849, + 1.9832, + 2.1202, + 2.2776, + 2.4286, + 2.5509, + 2.6357, + 2.6825, + 2.6937, + 2.6701, + 2.6127, + 2.5223, + 2.4011, + 2.2574, + 2.108, + 1.9774, + 1.8841, + 1.8278, + 1.7983, + 1.7856, + 1.7837 + ], + "Test Loss": [ + 1.7779, + 1.7799, + 1.7935, + 1.8252, + 1.8849, + 1.9832, + 2.1202, + 2.2776, + 2.4286, + 2.5509, + 2.6357, + 2.6825, + 2.6937, + 2.6701, + 2.6127, + 2.5223, + 2.4011, + 2.2574, + 2.108, + 1.9774, + 1.8841, + 1.8278, + 1.7983, + 1.7856, + 1.7837 + ], + "Val PPL": [ + 5.9177, + 5.9295, + 6.0106, + 6.2042, + 6.5855, + 7.2661, + 8.3328, + 9.7529, + 11.3425, + 12.8187, + 13.9533, + 14.6217, + 14.7857, + 14.4415, + 13.6361, + 12.4575, + 11.0355, + 9.5579, + 8.2316, + 7.2242, + 6.5802, + 6.2201, + 6.0395, + 5.9629, + 5.9519 + ], + "Test PPL": [ + 5.9177, + 5.9295, + 6.0106, + 6.2042, + 6.5855, + 7.2661, + 8.3328, + 9.7529, + 11.3425, + 12.8187, + 13.9533, + 14.6217, + 14.7857, + 14.4415, + 13.6361, + 12.4575, + 11.0355, + 9.5579, + 8.2316, + 7.2242, + 6.5802, + 6.2201, + 6.0395, + 5.9629, + 5.9519 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.7779, + 1.7788, + 1.7849, + 1.7976, + 1.8181, + 1.8476, + 1.8865, + 1.9328, + 1.9827, + 2.0313, + 2.0739, + 2.1056, + 2.1225, + 2.1221, + 2.1043, + 2.0709, + 2.0256, + 1.9731, + 1.9199, + 1.8724, + 1.8347, + 1.8084, + 1.7924, + 1.7847, + 1.7837 + ], + "Test Loss": [ + 1.7779, + 1.7788, + 1.7849, + 1.7976, + 1.8181, + 1.8476, + 1.8865, + 1.9328, + 1.9827, + 2.0313, + 2.0739, + 2.1056, + 2.1225, + 2.1221, + 2.1043, + 2.0709, + 2.0256, + 1.9731, + 1.9199, + 1.8724, + 1.8347, + 1.8084, + 1.7924, + 1.7847, + 1.7837 + ], + "Val PPL": [ + 5.9177, + 5.923, + 5.9593, + 6.0349, + 6.16, + 6.3449, + 6.5965, + 6.909, + 7.2623, + 7.6241, + 7.956, + 8.2116, + 8.3518, + 8.3489, + 8.2012, + 7.9318, + 7.5807, + 7.1926, + 6.8205, + 6.504, + 6.2631, + 6.101, + 6.004, + 5.9578, + 5.9518 + ], + "Test PPL": [ + 5.9177, + 5.923, + 5.9593, + 6.0349, + 6.16, + 6.3449, + 6.5965, + 6.909, + 7.2623, + 7.6241, + 7.956, + 8.2116, + 8.3518, + 8.3489, + 8.2012, + 7.9318, + 7.5807, + 7.1926, + 6.8205, + 6.504, + 6.2631, + 6.101, + 6.004, + 5.9578, + 5.9518 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.7779, + 1.7786, + 1.7837, + 1.7946, + 1.8118, + 1.8364, + 1.8688, + 1.9084, + 1.9533, + 2.0, + 2.0446, + 2.0826, + 2.1102, + 2.1237, + 2.1205, + 2.1, + 2.0632, + 2.0138, + 1.9577, + 1.9021, + 1.8543, + 1.8191, + 1.7968, + 1.7857, + 1.7837 + ], + "Test Loss": [ + 1.7779, + 1.7786, + 1.7837, + 1.7946, + 1.8118, + 1.8364, + 1.8688, + 1.9084, + 1.9533, + 2.0, + 2.0446, + 2.0826, + 2.1102, + 2.1237, + 2.1205, + 2.1, + 2.0632, + 2.0138, + 1.9577, + 1.9021, + 1.8543, + 1.8191, + 1.7968, + 1.7857, + 1.7837 + ], + "Val PPL": [ + 5.9177, + 5.9214, + 5.9521, + 6.0169, + 6.1214, + 6.2741, + 6.4803, + 6.7424, + 7.0517, + 7.389, + 7.7262, + 8.0257, + 8.25, + 8.3619, + 8.3354, + 8.1664, + 7.8712, + 7.4916, + 7.0828, + 6.6997, + 6.387, + 6.1662, + 6.0304, + 5.964, + 5.9516 + ], + "Test PPL": [ + 5.9177, + 5.9214, + 5.9521, + 6.0169, + 6.1214, + 6.2741, + 6.4803, + 6.7424, + 7.0517, + 7.389, + 7.7262, + 8.0257, + 8.25, + 8.3619, + 8.3354, + 8.1664, + 7.8712, + 7.4916, + 7.0828, + 6.6997, + 6.387, + 6.1662, + 6.0304, + 5.964, + 5.9516 + ] + } +} \ No newline at end of file diff --git a/results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json b/results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..8b7cb87daed76f0e73cd56f31244fdc25bb2e4a3 --- /dev/null +++ b/results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.7779, + 1.78, + 1.7922, + 1.8203, + 1.8739, + 1.9653, + 2.0983, + 2.2572, + 2.4126, + 2.5411, + 2.6319, + 2.6834, + 2.6965, + 2.6721, + 2.6109, + 2.5142, + 2.3875, + 2.2426, + 2.099, + 1.978, + 1.8917, + 1.8378, + 1.8075, + 1.7928, + 1.7892 + ], + "Test Loss": [ + 1.7779, + 1.78, + 1.7922, + 1.8203, + 1.8739, + 1.9653, + 2.0983, + 2.2572, + 2.4126, + 2.5411, + 2.6319, + 2.6834, + 2.6965, + 2.6721, + 2.6109, + 2.5142, + 2.3875, + 2.2426, + 2.099, + 1.978, + 1.8917, + 1.8378, + 1.8075, + 1.7928, + 1.7892 + ], + "Val PPL": [ + 5.9177, + 5.9298, + 6.0028, + 6.1738, + 6.5138, + 7.1367, + 8.1526, + 9.5564, + 11.1635, + 12.6934, + 13.9006, + 14.6347, + 14.8276, + 14.4699, + 13.6112, + 12.3563, + 10.8858, + 9.4174, + 8.1584, + 7.2281, + 6.6304, + 6.2826, + 6.0949, + 6.006, + 5.9845 + ], + "Test PPL": [ + 5.9177, + 5.9298, + 6.0028, + 6.1738, + 6.5138, + 7.1367, + 8.1526, + 9.5564, + 11.1635, + 12.6934, + 13.9006, + 14.6347, + 14.8276, + 14.4699, + 13.6112, + 12.3563, + 10.8858, + 9.4174, + 8.1584, + 7.2281, + 6.6304, + 6.2826, + 6.0949, + 6.006, + 5.9845 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.7779, + 1.7791, + 1.7859, + 1.7996, + 1.8218, + 1.8543, + 1.8972, + 1.9484, + 2.0025, + 2.0534, + 2.0949, + 2.122, + 2.1319, + 2.1235, + 2.0981, + 2.0587, + 2.0105, + 1.9588, + 1.9099, + 1.868, + 1.8355, + 1.8125, + 1.7981, + 1.7907, + 1.7892 + ], + "Test Loss": [ + 1.7779, + 1.7791, + 1.7859, + 1.7996, + 1.8218, + 1.8543, + 1.8972, + 1.9484, + 2.0025, + 2.0534, + 2.0949, + 2.122, + 2.1319, + 2.1235, + 2.0981, + 2.0587, + 2.0105, + 1.9588, + 1.9099, + 1.868, + 1.8355, + 1.8125, + 1.7981, + 1.7907, + 1.7892 + ], + "Val PPL": [ + 5.9177, + 5.9247, + 5.9649, + 6.0471, + 6.1829, + 6.3871, + 6.6673, + 7.0172, + 7.4074, + 7.7944, + 8.1249, + 8.348, + 8.4307, + 8.3605, + 8.1507, + 7.8358, + 7.4667, + 7.0911, + 6.7526, + 6.4751, + 6.2684, + 6.1257, + 6.0379, + 5.9936, + 5.9845 + ], + "Test PPL": [ + 5.9177, + 5.9247, + 5.9649, + 6.0471, + 6.1829, + 6.3871, + 6.6673, + 7.0172, + 7.4074, + 7.7944, + 8.1249, + 8.348, + 8.4307, + 8.3605, + 8.1507, + 7.8358, + 7.4667, + 7.0911, + 6.7526, + 6.4751, + 6.2684, + 6.1257, + 6.0379, + 5.9936, + 5.9845 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.7779, + 1.7789, + 1.7845, + 1.796, + 1.8147, + 1.8422, + 1.879, + 1.9246, + 1.9759, + 2.0282, + 2.0759, + 2.1136, + 2.1368, + 2.1421, + 2.1289, + 2.0985, + 2.0537, + 2.0004, + 1.9445, + 1.8929, + 1.8508, + 1.8203, + 1.8011, + 1.7913, + 1.7892 + ], + "Test Loss": [ + 1.7779, + 1.7789, + 1.7845, + 1.796, + 1.8147, + 1.8422, + 1.879, + 1.9246, + 1.9759, + 2.0282, + 2.0759, + 2.1136, + 2.1368, + 2.1421, + 2.1289, + 2.0985, + 2.0537, + 2.0004, + 1.9445, + 1.8929, + 1.8508, + 1.8203, + 1.8011, + 1.7913, + 1.7892 + ], + "Val PPL": [ + 5.9177, + 5.9232, + 5.9566, + 6.0254, + 6.1391, + 6.3102, + 6.5469, + 6.8524, + 7.213, + 7.6004, + 7.9715, + 8.2783, + 8.4725, + 8.5175, + 8.4053, + 8.1538, + 7.7969, + 7.3921, + 6.9904, + 6.6384, + 6.3646, + 6.1736, + 6.0565, + 5.9971, + 5.9844 + ], + "Test PPL": [ + 5.9177, + 5.9232, + 5.9566, + 6.0254, + 6.1391, + 6.3102, + 6.5469, + 6.8524, + 7.213, + 7.6004, + 7.9715, + 8.2783, + 8.4725, + 8.5175, + 8.4053, + 8.1538, + 7.7969, + 7.3921, + 6.9904, + 6.6384, + 6.3646, + 6.1736, + 6.0565, + 5.9971, + 5.9844 + ] + } +} \ No newline at end of file diff --git a/results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json b/results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..35847e42522a3dcf6c1f936d1a858736a88a56db --- /dev/null +++ b/results/text8/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.7837, + 1.7853, + 1.797, + 1.8248, + 1.8794, + 1.9733, + 2.1078, + 2.265, + 2.4178, + 2.546, + 2.6394, + 2.6937, + 2.7084, + 2.6843, + 2.6231, + 2.5269, + 2.3993, + 2.2507, + 2.1004, + 1.9734, + 1.8846, + 1.832, + 1.8041, + 1.7917, + 1.7892 + ], + "Test Loss": [ + 1.7837, + 1.7853, + 1.797, + 1.8248, + 1.8794, + 1.9733, + 2.1078, + 2.265, + 2.4178, + 2.546, + 2.6394, + 2.6937, + 2.7084, + 2.6843, + 2.6231, + 2.5269, + 2.3993, + 2.2507, + 2.1004, + 1.9734, + 1.8846, + 1.832, + 1.8041, + 1.7917, + 1.7892 + ], + "Val PPL": [ + 5.9519, + 5.9616, + 6.0315, + 6.2017, + 6.5499, + 7.1941, + 8.2301, + 9.6313, + 11.221, + 12.7554, + 14.0054, + 14.7868, + 15.0053, + 14.6479, + 13.7789, + 12.5145, + 11.0157, + 9.4942, + 8.1698, + 7.1948, + 6.5837, + 6.2464, + 6.0747, + 5.9999, + 5.9845 + ], + "Test PPL": [ + 5.9519, + 5.9616, + 6.0315, + 6.2017, + 6.5499, + 7.1941, + 8.2301, + 9.6313, + 11.221, + 12.7554, + 14.0054, + 14.7868, + 15.0053, + 14.6479, + 13.7789, + 12.5145, + 11.0157, + 9.4942, + 8.1698, + 7.1948, + 6.5837, + 6.2464, + 6.0747, + 5.9999, + 5.9845 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.7837, + 1.7846, + 1.7911, + 1.8048, + 1.8273, + 1.8595, + 1.9013, + 1.9502, + 2.0021, + 2.0507, + 2.0899, + 2.1147, + 2.1222, + 2.1119, + 2.0857, + 2.0471, + 2.0006, + 1.9516, + 1.905, + 1.8649, + 1.8337, + 1.8116, + 1.7977, + 1.7908, + 1.7892 + ], + "Test Loss": [ + 1.7837, + 1.7846, + 1.7911, + 1.8048, + 1.8273, + 1.8595, + 1.9013, + 1.9502, + 2.0021, + 2.0507, + 2.0899, + 2.1147, + 2.1222, + 2.1119, + 2.0857, + 2.0471, + 2.0006, + 1.9516, + 1.905, + 1.8649, + 1.8337, + 1.8116, + 1.7977, + 1.7908, + 1.7892 + ], + "Val PPL": [ + 5.9519, + 5.9573, + 5.9958, + 6.0785, + 6.2168, + 6.4204, + 6.6946, + 7.0302, + 7.4045, + 7.7731, + 8.0844, + 8.2874, + 8.3496, + 8.2639, + 8.0503, + 7.7454, + 7.3935, + 7.0399, + 6.7191, + 6.4555, + 6.2567, + 6.12, + 6.0357, + 5.9941, + 5.9847 + ], + "Test PPL": [ + 5.9519, + 5.9573, + 5.9958, + 6.0785, + 6.2168, + 6.4204, + 6.6946, + 7.0302, + 7.4045, + 7.7731, + 8.0844, + 8.2874, + 8.3496, + 8.2639, + 8.0503, + 7.7454, + 7.3935, + 7.0399, + 6.7191, + 6.4555, + 6.2567, + 6.12, + 6.0357, + 5.9941, + 5.9847 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.7837, + 1.7842, + 1.7895, + 1.8011, + 1.8205, + 1.8484, + 1.885, + 1.9291, + 1.9784, + 2.0292, + 2.0762, + 2.1137, + 2.1369, + 2.1424, + 2.129, + 2.0979, + 2.0523, + 1.9977, + 1.9414, + 1.8903, + 1.849, + 1.8196, + 1.8009, + 1.7915, + 1.7892 + ], + "Test Loss": [ + 1.7837, + 1.7842, + 1.7895, + 1.8011, + 1.8205, + 1.8484, + 1.885, + 1.9291, + 1.9784, + 2.0292, + 2.0762, + 2.1137, + 2.1369, + 2.1424, + 2.129, + 2.0979, + 2.0523, + 1.9977, + 1.9414, + 1.8903, + 1.849, + 1.8196, + 1.8009, + 1.7915, + 1.7892 + ], + "Val PPL": [ + 5.9519, + 5.9547, + 5.9863, + 6.0564, + 6.175, + 6.3494, + 6.5861, + 6.8831, + 7.2314, + 7.6081, + 7.9737, + 8.279, + 8.4733, + 8.5199, + 8.4061, + 8.1489, + 7.7855, + 7.3722, + 6.9684, + 6.6211, + 6.3534, + 6.1692, + 6.0553, + 5.9983, + 5.9845 + ], + "Test PPL": [ + 5.9519, + 5.9547, + 5.9863, + 6.0564, + 6.175, + 6.3494, + 6.5861, + 6.8831, + 7.2314, + 7.6081, + 7.9737, + 8.279, + 8.4733, + 8.5199, + 8.4061, + 8.1489, + 7.7855, + 7.3722, + 6.9684, + 6.6211, + 6.3534, + 6.1692, + 6.0553, + 5.9983, + 5.9845 + ] + } +} \ No newline at end of file diff --git a/results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json b/results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..26a41beb47898445f19ec66b346d6d9f42bdbbce --- /dev/null +++ b/results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.9355, + 1.9374, + 1.9481, + 1.9763, + 2.0349, + 2.1222, + 2.23, + 2.3502, + 2.4696, + 2.5747, + 2.6567, + 2.7108, + 2.7357, + 2.7313, + 2.6974, + 2.6351, + 2.5459, + 2.4359, + 2.3124, + 2.1866, + 2.0723, + 1.9859, + 1.9485, + 1.9379, + 1.9366 + ], + "Test Loss": [ + 1.9355, + 1.9374, + 1.9481, + 1.9763, + 2.0349, + 2.1222, + 2.23, + 2.3502, + 2.4696, + 2.5747, + 2.6567, + 2.7108, + 2.7357, + 2.7313, + 2.6974, + 2.6351, + 2.5459, + 2.4359, + 2.3124, + 2.1866, + 2.0723, + 1.9859, + 1.9485, + 1.9379, + 1.9366 + ], + "Val PPL": [ + 6.9276, + 6.9405, + 7.0154, + 7.2161, + 7.6514, + 8.3494, + 9.2997, + 10.4875, + 11.8176, + 13.1278, + 14.2497, + 15.0409, + 15.4202, + 15.3529, + 14.8413, + 13.9443, + 12.7548, + 11.4266, + 10.0984, + 8.9046, + 7.9429, + 7.2858, + 7.0184, + 6.9441, + 6.9353 + ], + "Test PPL": [ + 6.9276, + 6.9405, + 7.0154, + 7.2161, + 7.6514, + 8.3494, + 9.2997, + 10.4875, + 11.8176, + 13.1278, + 14.2497, + 15.0409, + 15.4202, + 15.3529, + 14.8413, + 13.9443, + 12.7548, + 11.4266, + 10.0984, + 8.9046, + 7.9429, + 7.2858, + 7.0184, + 6.9441, + 6.9353 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.9355, + 1.9399, + 1.9495, + 1.9652, + 1.9877, + 2.0176, + 2.0551, + 2.0991, + 2.1453, + 2.1884, + 2.2226, + 2.2442, + 2.2518, + 2.245, + 2.2246, + 2.1921, + 2.1502, + 2.1036, + 2.0577, + 2.0173, + 1.9853, + 1.9625, + 1.9481, + 1.9399, + 1.9366 + ], + "Test Loss": [ + 1.9355, + 1.9399, + 1.9495, + 1.9652, + 1.9877, + 2.0176, + 2.0551, + 2.0991, + 2.1453, + 2.1884, + 2.2226, + 2.2442, + 2.2518, + 2.245, + 2.2246, + 2.1921, + 2.1502, + 2.1036, + 2.0577, + 2.0173, + 1.9853, + 1.9625, + 1.9481, + 1.9399, + 1.9366 + ], + "Val PPL": [ + 6.9276, + 6.9579, + 7.0254, + 7.1362, + 7.2986, + 7.5206, + 7.8076, + 8.1587, + 8.545, + 8.921, + 9.2315, + 9.4326, + 9.5048, + 9.4407, + 9.2501, + 8.954, + 8.5868, + 8.1959, + 7.8283, + 7.5177, + 7.2809, + 7.1174, + 7.0152, + 6.9577, + 6.9355 + ], + "Test PPL": [ + 6.9276, + 6.9579, + 7.0254, + 7.1362, + 7.2986, + 7.5206, + 7.8076, + 8.1587, + 8.545, + 8.921, + 9.2315, + 9.4326, + 9.5048, + 9.4407, + 9.2501, + 8.954, + 8.5868, + 8.1959, + 7.8283, + 7.5177, + 7.2809, + 7.1174, + 7.0152, + 6.9577, + 6.9355 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.9355, + 1.9388, + 1.9457, + 1.957, + 1.9735, + 1.9962, + 2.0262, + 2.064, + 2.1079, + 2.1535, + 2.1961, + 2.2307, + 2.2535, + 2.2619, + 2.2558, + 2.2342, + 2.1993, + 2.1549, + 2.105, + 2.0542, + 2.0087, + 1.9747, + 1.953, + 1.941, + 1.9366 + ], + "Test Loss": [ + 1.9355, + 1.9388, + 1.9457, + 1.957, + 1.9735, + 1.9962, + 2.0262, + 2.064, + 2.1079, + 2.1535, + 2.1961, + 2.2307, + 2.2535, + 2.2619, + 2.2558, + 2.2342, + 2.1993, + 2.1549, + 2.105, + 2.0542, + 2.0087, + 1.9747, + 1.953, + 1.941, + 1.9366 + ], + "Val PPL": [ + 6.9276, + 6.9502, + 6.9988, + 7.0782, + 7.1961, + 7.3613, + 7.5855, + 7.8776, + 8.2308, + 8.6152, + 8.9899, + 9.3068, + 9.5208, + 9.6016, + 9.5426, + 9.3389, + 9.0187, + 8.6272, + 8.2068, + 7.8003, + 7.4537, + 7.2044, + 7.0499, + 6.9655, + 6.935 + ], + "Test PPL": [ + 6.9276, + 6.9502, + 6.9988, + 7.0782, + 7.1961, + 7.3613, + 7.5855, + 7.8776, + 8.2308, + 8.6152, + 8.9899, + 9.3068, + 9.5208, + 9.6016, + 9.5426, + 9.3389, + 9.0187, + 8.6272, + 8.2068, + 7.8003, + 7.4537, + 7.2044, + 7.0499, + 6.9655, + 6.935 + ] + } +} \ No newline at end of file diff --git a/results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json b/results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..604d08f168d2fa6415814a04d052430c89f4edf7 --- /dev/null +++ b/results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.9355, + 1.9369, + 1.9464, + 1.9731, + 2.0241, + 2.0997, + 2.1982, + 2.3103, + 2.4217, + 2.5191, + 2.5931, + 2.6379, + 2.6512, + 2.6324, + 2.5821, + 2.5044, + 2.406, + 2.2971, + 2.1897, + 2.0952, + 2.0194, + 1.9661, + 1.9379, + 1.9285, + 1.9282 + ], + "Test Loss": [ + 1.9355, + 1.9369, + 1.9464, + 1.9731, + 2.0241, + 2.0997, + 2.1982, + 2.3103, + 2.4217, + 2.5191, + 2.5931, + 2.6379, + 2.6512, + 2.6324, + 2.5821, + 2.5044, + 2.406, + 2.2971, + 2.1897, + 2.0952, + 2.0194, + 1.9661, + 1.9379, + 1.9285, + 1.9282 + ], + "Val PPL": [ + 6.9276, + 6.9371, + 7.0036, + 7.1931, + 7.569, + 8.1635, + 9.0084, + 10.0774, + 11.2655, + 12.4171, + 13.3706, + 13.9832, + 14.1717, + 13.9076, + 13.2254, + 12.2358, + 11.0895, + 9.9454, + 8.9327, + 8.1273, + 7.5336, + 7.1427, + 6.9444, + 6.8795, + 6.8769 + ], + "Test PPL": [ + 6.9276, + 6.9371, + 7.0036, + 7.1931, + 7.569, + 8.1635, + 9.0084, + 10.0774, + 11.2655, + 12.4171, + 13.3706, + 13.9832, + 14.1717, + 13.9076, + 13.2254, + 12.2358, + 11.0895, + 9.9454, + 8.9327, + 8.1273, + 7.5336, + 7.1427, + 6.9444, + 6.8795, + 6.8769 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.9355, + 1.9378, + 1.9448, + 1.9577, + 1.9777, + 2.006, + 2.0438, + 2.091, + 2.1427, + 2.1912, + 2.2286, + 2.2501, + 2.2548, + 2.2446, + 2.222, + 2.189, + 2.148, + 2.1016, + 2.0545, + 2.0118, + 1.9775, + 1.9531, + 1.9379, + 1.93, + 1.9282 + ], + "Test Loss": [ + 1.9355, + 1.9378, + 1.9448, + 1.9577, + 1.9777, + 2.006, + 2.0438, + 2.091, + 2.1427, + 2.1912, + 2.2286, + 2.2501, + 2.2548, + 2.2446, + 2.222, + 2.189, + 2.148, + 2.1016, + 2.0545, + 2.0118, + 1.9775, + 1.9531, + 1.9379, + 1.93, + 1.9282 + ], + "Val PPL": [ + 6.9276, + 6.9434, + 6.9922, + 7.0831, + 7.2262, + 7.4336, + 7.7201, + 8.0932, + 8.5227, + 8.946, + 9.287, + 9.4885, + 9.533, + 9.437, + 9.2257, + 8.9263, + 8.5675, + 8.1795, + 7.8032, + 7.4765, + 7.2247, + 7.0507, + 6.9439, + 6.8894, + 6.8769 + ], + "Test PPL": [ + 6.9276, + 6.9434, + 6.9922, + 7.0831, + 7.2262, + 7.4336, + 7.7201, + 8.0932, + 8.5227, + 8.946, + 9.287, + 9.4885, + 9.533, + 9.437, + 9.2257, + 8.9263, + 8.5675, + 8.1795, + 7.8032, + 7.4765, + 7.2247, + 7.0507, + 6.9439, + 6.8894, + 6.8769 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.9355, + 1.9374, + 1.9426, + 1.9517, + 1.9656, + 1.9851, + 2.0115, + 2.0458, + 2.0879, + 2.1358, + 2.1833, + 2.2236, + 2.2513, + 2.2631, + 2.2588, + 2.2403, + 2.2092, + 2.1672, + 2.1168, + 2.0632, + 2.0124, + 1.9719, + 1.9452, + 1.9317, + 1.9281 + ], + "Test Loss": [ + 1.9355, + 1.9374, + 1.9426, + 1.9517, + 1.9656, + 1.9851, + 2.0115, + 2.0458, + 2.0879, + 2.1358, + 2.1833, + 2.2236, + 2.2513, + 2.2631, + 2.2588, + 2.2403, + 2.2092, + 2.1672, + 2.1168, + 2.0632, + 2.0124, + 1.9719, + 1.9452, + 1.9317, + 1.9281 + ], + "Val PPL": [ + 6.9276, + 6.941, + 6.9768, + 7.0404, + 7.139, + 7.2795, + 7.4743, + 7.7351, + 8.0679, + 8.4636, + 8.876, + 9.2408, + 9.4998, + 9.6131, + 9.5717, + 9.396, + 9.1085, + 8.7334, + 8.3044, + 7.8709, + 7.4815, + 7.1841, + 6.9953, + 6.901, + 6.8767 + ], + "Test PPL": [ + 6.9276, + 6.941, + 6.9768, + 7.0404, + 7.139, + 7.2795, + 7.4743, + 7.7351, + 8.0679, + 8.4636, + 8.876, + 9.2408, + 9.4998, + 9.6131, + 9.5717, + 9.396, + 9.1085, + 8.7334, + 8.3044, + 7.8709, + 7.4815, + 7.1841, + 6.9953, + 6.901, + 6.8767 + ] + } +} \ No newline at end of file diff --git a/results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json b/results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..c5addf900f8b4c3a3e46059e3a6e589d1adc11ec --- /dev/null +++ b/results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.9366, + 1.9387, + 1.9489, + 1.9761, + 2.0401, + 2.1493, + 2.2862, + 2.4281, + 2.5589, + 2.6624, + 2.7331, + 2.7684, + 2.7696, + 2.7371, + 2.6745, + 2.5866, + 2.478, + 2.3572, + 2.2352, + 2.1249, + 2.0373, + 1.9769, + 1.9431, + 1.9301, + 1.9282 + ], + "Test Loss": [ + 1.9366, + 1.9387, + 1.9489, + 1.9761, + 2.0401, + 2.1493, + 2.2862, + 2.4281, + 2.5589, + 2.6624, + 2.7331, + 2.7684, + 2.7696, + 2.7371, + 2.6745, + 2.5866, + 2.478, + 2.3572, + 2.2352, + 2.1249, + 2.0373, + 1.9769, + 1.9431, + 1.9301, + 1.9282 + ], + "Val PPL": [ + 6.9353, + 6.9495, + 7.021, + 7.2146, + 7.691, + 8.5786, + 9.8371, + 11.3377, + 12.921, + 14.3308, + 15.38, + 15.9332, + 15.9515, + 15.4427, + 14.5046, + 13.2844, + 11.9173, + 10.5609, + 9.3483, + 8.372, + 7.6697, + 7.2204, + 6.9807, + 6.8899, + 6.8769 + ], + "Test PPL": [ + 6.9353, + 6.9495, + 7.021, + 7.2146, + 7.691, + 8.5786, + 9.8371, + 11.3377, + 12.921, + 14.3308, + 15.38, + 15.9332, + 15.9515, + 15.4427, + 14.5046, + 13.2844, + 11.9173, + 10.5609, + 9.3483, + 8.372, + 7.6697, + 7.2204, + 6.9807, + 6.8899, + 6.8769 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.9366, + 1.9388, + 1.9461, + 1.9606, + 1.9853, + 2.024, + 2.0723, + 2.1193, + 2.1644, + 2.2063, + 2.2412, + 2.264, + 2.2695, + 2.2567, + 2.2297, + 2.194, + 2.1537, + 2.1116, + 2.0697, + 2.0296, + 1.9931, + 1.9618, + 1.9405, + 1.9306, + 1.9282 + ], + "Test Loss": [ + 1.9366, + 1.9388, + 1.9461, + 1.9606, + 1.9853, + 2.024, + 2.0723, + 2.1193, + 2.1644, + 2.2063, + 2.2412, + 2.264, + 2.2695, + 2.2567, + 2.2297, + 2.194, + 2.1537, + 2.1116, + 2.0697, + 2.0296, + 1.9931, + 1.9618, + 1.9405, + 1.9306, + 1.9282 + ], + "Val PPL": [ + 6.9353, + 6.9504, + 7.0016, + 7.1039, + 7.2811, + 7.5684, + 7.9432, + 8.3253, + 8.7097, + 9.082, + 9.4048, + 9.6212, + 9.6742, + 9.5512, + 9.2974, + 8.9714, + 8.6164, + 8.2613, + 7.9222, + 7.611, + 7.3383, + 7.1121, + 6.9625, + 6.8933, + 6.8769 + ], + "Test PPL": [ + 6.9353, + 6.9504, + 7.0016, + 7.1039, + 7.2811, + 7.5684, + 7.9432, + 8.3253, + 8.7097, + 9.082, + 9.4048, + 9.6212, + 9.6742, + 9.5512, + 9.2974, + 8.9714, + 8.6164, + 8.2613, + 7.9222, + 7.611, + 7.3383, + 7.1121, + 6.9625, + 6.8933, + 6.8769 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.9366, + 1.9386, + 1.9459, + 1.9673, + 2.0073, + 2.0557, + 2.1162, + 2.1868, + 2.2548, + 2.3108, + 2.3516, + 2.3771, + 2.3861, + 2.3748, + 2.3427, + 2.2926, + 2.2355, + 2.1772, + 2.1203, + 2.0652, + 2.0151, + 1.9752, + 1.9481, + 1.9333, + 1.9282 + ], + "Test Loss": [ + 1.9366, + 1.9386, + 1.9459, + 1.9673, + 2.0073, + 2.0557, + 2.1162, + 2.1868, + 2.2548, + 2.3108, + 2.3516, + 2.3771, + 2.3861, + 2.3748, + 2.3427, + 2.2926, + 2.2355, + 2.1772, + 2.1203, + 2.0652, + 2.0151, + 1.9752, + 1.9481, + 1.9333, + 1.9282 + ], + "Val PPL": [ + 6.9353, + 6.9487, + 7.0002, + 7.1512, + 7.4434, + 7.8123, + 8.2999, + 8.9063, + 9.5332, + 10.082, + 10.5027, + 10.774, + 10.8705, + 10.7487, + 10.4097, + 9.9004, + 9.351, + 8.822, + 8.334, + 7.8871, + 7.5016, + 7.2081, + 7.0151, + 6.9119, + 6.8769 + ], + "Test PPL": [ + 6.9353, + 6.9487, + 7.0002, + 7.1512, + 7.4434, + 7.8123, + 8.2999, + 8.9063, + 9.5332, + 10.082, + 10.5027, + 10.774, + 10.8705, + 10.7487, + 10.4097, + 9.9004, + 9.351, + 8.822, + 8.334, + 7.8871, + 7.5016, + 7.2081, + 7.0151, + 6.9119, + 6.8769 + ] + } +} \ No newline at end of file diff --git a/results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json b/results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..bada9152458f018bc4ca57a3d85c966697da3a3e --- /dev/null +++ b/results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.8373, + 1.8377, + 1.8461, + 1.867, + 1.9079, + 1.9776, + 2.0816, + 2.2138, + 2.3558, + 2.4856, + 2.5869, + 2.6513, + 2.6753, + 2.6585, + 2.6009, + 2.5062, + 2.3808, + 2.2395, + 2.1038, + 1.9934, + 1.9165, + 1.8694, + 1.8438, + 1.8327, + 1.8313 + ], + "Test Loss": [ + 1.8373, + 1.8377, + 1.8461, + 1.867, + 1.9079, + 1.9776, + 2.0816, + 2.2138, + 2.3558, + 2.4856, + 2.5869, + 2.6513, + 2.6753, + 2.6585, + 2.6009, + 2.5062, + 2.3808, + 2.2395, + 2.1038, + 1.9934, + 1.9165, + 1.8694, + 1.8438, + 1.8327, + 1.8313 + ], + "Val PPL": [ + 6.2794, + 6.2818, + 6.335, + 6.4689, + 6.7389, + 7.2257, + 8.0171, + 9.1508, + 10.5464, + 12.0082, + 13.2888, + 14.1721, + 14.5167, + 14.2743, + 13.4759, + 12.2578, + 10.8133, + 9.3884, + 8.1977, + 7.3405, + 6.7973, + 6.4846, + 6.3207, + 6.2509, + 6.2418 + ], + "Test PPL": [ + 6.2794, + 6.2818, + 6.335, + 6.4689, + 6.7389, + 7.2257, + 8.0171, + 9.1508, + 10.5464, + 12.0082, + 13.2888, + 14.1721, + 14.5167, + 14.2743, + 13.4759, + 12.2578, + 10.8133, + 9.3884, + 8.1977, + 7.3405, + 6.7973, + 6.4846, + 6.3207, + 6.2509, + 6.2418 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.8373, + 1.8405, + 1.8526, + 1.8762, + 1.9139, + 1.9667, + 2.0322, + 2.1047, + 2.1762, + 2.2391, + 2.2868, + 2.3151, + 2.3225, + 2.3082, + 2.2734, + 2.2241, + 2.1657, + 2.1037, + 2.0427, + 1.9859, + 1.9335, + 1.8781, + 1.8442, + 1.8332, + 1.8312 + ], + "Test Loss": [ + 1.8373, + 1.8405, + 1.8526, + 1.8762, + 1.9139, + 1.9667, + 2.0322, + 2.1047, + 2.1762, + 2.2391, + 2.2868, + 2.3151, + 2.3225, + 2.3082, + 2.2734, + 2.2241, + 2.1657, + 2.1037, + 2.0427, + 1.9859, + 1.9335, + 1.8781, + 1.8442, + 1.8332, + 1.8312 + ], + "Val PPL": [ + 6.2794, + 6.2997, + 6.3761, + 6.5285, + 6.7792, + 7.1467, + 7.6306, + 8.2045, + 8.8127, + 9.3848, + 9.8436, + 10.1263, + 10.2009, + 10.056, + 9.7125, + 9.2454, + 8.7206, + 8.1962, + 7.7116, + 7.2858, + 6.914, + 6.5412, + 6.3233, + 6.2539, + 6.2417 + ], + "Test PPL": [ + 6.2794, + 6.2997, + 6.3761, + 6.5285, + 6.7792, + 7.1467, + 7.6306, + 8.2045, + 8.8127, + 9.3848, + 9.8436, + 10.1263, + 10.2009, + 10.056, + 9.7125, + 9.2454, + 8.7206, + 8.1962, + 7.7116, + 7.2858, + 6.914, + 6.5412, + 6.3233, + 6.2539, + 6.2417 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.8373, + 1.8406, + 1.8513, + 1.8714, + 1.9032, + 1.9479, + 2.0041, + 2.0692, + 2.1392, + 2.2076, + 2.2686, + 2.3161, + 2.3461, + 2.3557, + 2.3404, + 2.3023, + 2.245, + 2.1759, + 2.1056, + 2.0384, + 1.977, + 1.9211, + 1.8583, + 1.8345, + 1.8313 + ], + "Test Loss": [ + 1.8373, + 1.8406, + 1.8513, + 1.8714, + 1.9032, + 1.9479, + 2.0041, + 2.0692, + 2.1392, + 2.2076, + 2.2686, + 2.3161, + 2.3461, + 2.3557, + 2.3404, + 2.3023, + 2.245, + 2.1759, + 2.1056, + 2.0384, + 1.977, + 1.9211, + 1.8583, + 1.8345, + 1.8313 + ], + "Val PPL": [ + 6.2794, + 6.3001, + 6.3678, + 6.4971, + 6.7075, + 7.0141, + 7.4197, + 7.9184, + 8.4923, + 9.0936, + 9.6655, + 10.1363, + 10.4445, + 10.5452, + 10.3852, + 9.9972, + 9.4405, + 8.8101, + 8.2122, + 7.6784, + 7.221, + 6.8287, + 6.4125, + 6.2622, + 6.2418 + ], + "Test PPL": [ + 6.2794, + 6.3001, + 6.3678, + 6.4971, + 6.7075, + 7.0141, + 7.4197, + 7.9184, + 8.4923, + 9.0936, + 9.6655, + 10.1363, + 10.4445, + 10.5452, + 10.3852, + 9.9972, + 9.4405, + 8.8101, + 8.2122, + 7.6784, + 7.221, + 6.8287, + 6.4125, + 6.2622, + 6.2418 + ] + } +} \ No newline at end of file diff --git a/results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json b/results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..592d65386e384a1f4989db95869bc38477d75b67 --- /dev/null +++ b/results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.8373, + 1.8384, + 1.8495, + 1.879, + 1.9408, + 2.0405, + 2.1685, + 2.3077, + 2.4386, + 2.5486, + 2.6288, + 2.6747, + 2.6832, + 2.6528, + 2.5857, + 2.4842, + 2.3576, + 2.2172, + 2.0828, + 1.9744, + 1.9015, + 1.859, + 1.8368, + 1.8271, + 1.8253 + ], + "Test Loss": [ + 1.8373, + 1.8384, + 1.8495, + 1.879, + 1.9408, + 2.0405, + 2.1685, + 2.3077, + 2.4386, + 2.5486, + 2.6288, + 2.6747, + 2.6832, + 2.6528, + 2.5857, + 2.4842, + 2.3576, + 2.2172, + 2.0828, + 1.9744, + 1.9015, + 1.859, + 1.8368, + 1.8271, + 1.8253 + ], + "Val PPL": [ + 6.2794, + 6.2865, + 6.3567, + 6.547, + 6.9644, + 7.6942, + 8.7451, + 10.0511, + 11.4571, + 12.789, + 13.8571, + 14.5073, + 14.6319, + 14.1935, + 13.2731, + 11.9916, + 10.5656, + 9.1814, + 8.0267, + 7.202, + 6.6957, + 6.417, + 6.2764, + 6.2158, + 6.2046 + ], + "Test PPL": [ + 6.2794, + 6.2865, + 6.3567, + 6.547, + 6.9644, + 7.6942, + 8.7451, + 10.0511, + 11.4571, + 12.789, + 13.8571, + 14.5073, + 14.6319, + 14.1935, + 13.2731, + 11.9916, + 10.5656, + 9.1814, + 8.0267, + 7.202, + 6.6957, + 6.417, + 6.2764, + 6.2158, + 6.2046 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.8373, + 1.84, + 1.8492, + 1.867, + 1.8955, + 1.9357, + 1.9868, + 2.045, + 2.1053, + 2.1626, + 2.2114, + 2.2457, + 2.2603, + 2.2539, + 2.2275, + 2.1858, + 2.1331, + 2.0749, + 2.0146, + 1.9573, + 1.9058, + 1.8632, + 1.8387, + 1.8283, + 1.8253 + ], + "Test Loss": [ + 1.8373, + 1.84, + 1.8492, + 1.867, + 1.8955, + 1.9357, + 1.9868, + 2.045, + 2.1053, + 2.1626, + 2.2114, + 2.2457, + 2.2603, + 2.2539, + 2.2275, + 2.1858, + 2.1331, + 2.0749, + 2.0146, + 1.9573, + 1.9058, + 1.8632, + 1.8387, + 1.8283, + 1.8253 + ], + "Val PPL": [ + 6.2794, + 6.2965, + 6.3545, + 6.4689, + 6.6557, + 6.929, + 7.2919, + 7.7289, + 8.2092, + 8.6933, + 9.1282, + 9.4469, + 9.5862, + 9.5244, + 9.2766, + 8.8981, + 8.4411, + 7.9635, + 7.4978, + 7.0805, + 6.7247, + 6.4443, + 6.2886, + 6.2233, + 6.2047 + ], + "Test PPL": [ + 6.2794, + 6.2965, + 6.3545, + 6.4689, + 6.6557, + 6.929, + 7.2919, + 7.7289, + 8.2092, + 8.6933, + 9.1282, + 9.4469, + 9.5862, + 9.5244, + 9.2766, + 8.8981, + 8.4411, + 7.9635, + 7.4978, + 7.0805, + 6.7247, + 6.4443, + 6.2886, + 6.2233, + 6.2047 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.8373, + 1.84, + 1.8652, + 1.9105, + 1.9527, + 2.0061, + 2.0683, + 2.1337, + 2.1992, + 2.2592, + 2.3092, + 2.3446, + 2.3609, + 2.3565, + 2.332, + 2.2898, + 2.2331, + 2.1663, + 2.0931, + 2.0201, + 1.9515, + 1.893, + 1.8519, + 1.8312, + 1.8253 + ], + "Test Loss": [ + 1.8373, + 1.84, + 1.8652, + 1.9105, + 1.9527, + 2.0061, + 2.0683, + 2.1337, + 2.1992, + 2.2592, + 2.3092, + 2.3446, + 2.3609, + 2.3565, + 2.332, + 2.2898, + 2.2331, + 2.1663, + 2.0931, + 2.0201, + 1.9515, + 1.893, + 1.8519, + 1.8312, + 1.8253 + ], + "Val PPL": [ + 6.2794, + 6.2966, + 6.4572, + 6.7564, + 7.0476, + 7.4345, + 7.9114, + 8.4462, + 9.0182, + 9.575, + 10.0659, + 10.4296, + 10.6006, + 10.5542, + 10.2984, + 9.8729, + 9.329, + 8.7259, + 8.11, + 7.5388, + 7.0392, + 6.6396, + 6.372, + 6.2415, + 6.2046 + ], + "Test PPL": [ + 6.2794, + 6.2966, + 6.4572, + 6.7564, + 7.0476, + 7.4345, + 7.9114, + 8.4462, + 9.0182, + 9.575, + 10.0659, + 10.4296, + 10.6006, + 10.5542, + 10.2984, + 9.8729, + 9.329, + 8.7259, + 8.11, + 7.5388, + 7.0392, + 6.6396, + 6.372, + 6.2415, + 6.2046 + ] + } +} \ No newline at end of file diff --git a/results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json b/results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..b58bb7a7aaa3f143356fa1e388ab606f3624b1ad --- /dev/null +++ b/results/text8/[finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 1.8313, + 1.8341, + 1.8557, + 1.9033, + 1.9863, + 2.0971, + 2.2204, + 2.3456, + 2.4635, + 2.565, + 2.6432, + 2.6946, + 2.7161, + 2.7084, + 2.6735, + 2.6141, + 2.5347, + 2.4403, + 2.3373, + 2.2282, + 2.1072, + 1.9755, + 1.8585, + 1.8265, + 1.8253 + ], + "Test Loss": [ + 1.8313, + 1.8341, + 1.8557, + 1.9033, + 1.9863, + 2.0971, + 2.2204, + 2.3456, + 2.4635, + 2.565, + 2.6432, + 2.6946, + 2.7161, + 2.7084, + 2.6735, + 2.6141, + 2.5347, + 2.4403, + 2.3373, + 2.2282, + 2.1072, + 1.9755, + 1.8585, + 1.8265, + 1.8253 + ], + "Val PPL": [ + 6.2418, + 6.2596, + 6.396, + 6.7079, + 7.2884, + 8.1427, + 9.2113, + 10.4398, + 11.7454, + 13.0005, + 14.0585, + 14.7999, + 15.121, + 15.0054, + 14.4907, + 13.6556, + 12.6132, + 11.476, + 10.3528, + 9.2827, + 8.2251, + 7.2102, + 6.4142, + 6.2122, + 6.2046 + ], + "Test PPL": [ + 6.2418, + 6.2596, + 6.396, + 6.7079, + 7.2884, + 8.1427, + 9.2113, + 10.4398, + 11.7454, + 13.0005, + 14.0585, + 14.7999, + 15.121, + 15.0054, + 14.4907, + 13.6556, + 12.6132, + 11.476, + 10.3528, + 9.2827, + 8.2251, + 7.2102, + 6.4142, + 6.2122, + 6.2046 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 1.8313, + 1.8332, + 1.8418, + 1.8586, + 1.8853, + 1.9226, + 1.9694, + 2.0227, + 2.0781, + 2.1284, + 2.1663, + 2.1873, + 2.1902, + 2.1756, + 2.1464, + 2.106, + 2.0586, + 2.0085, + 1.961, + 1.919, + 1.8837, + 1.8557, + 1.8377, + 1.8283, + 1.8253 + ], + "Test Loss": [ + 1.8313, + 1.8332, + 1.8418, + 1.8586, + 1.8853, + 1.9226, + 1.9694, + 2.0227, + 2.0781, + 2.1284, + 2.1663, + 2.1873, + 2.1902, + 2.1756, + 2.1464, + 2.106, + 2.0586, + 2.0085, + 1.961, + 1.919, + 1.8837, + 1.8557, + 1.8377, + 1.8283, + 1.8253 + ], + "Val PPL": [ + 6.2418, + 6.2537, + 6.308, + 6.415, + 6.5884, + 6.8386, + 7.1662, + 7.5588, + 7.9894, + 8.4013, + 8.7264, + 8.9108, + 8.9371, + 8.8073, + 8.5544, + 8.215, + 7.8348, + 7.4522, + 7.1063, + 6.8145, + 6.5775, + 6.3964, + 6.2819, + 6.2235, + 6.2045 + ], + "Test PPL": [ + 6.2418, + 6.2537, + 6.308, + 6.415, + 6.5884, + 6.8386, + 7.1662, + 7.5588, + 7.9894, + 8.4013, + 8.7264, + 8.9108, + 8.9371, + 8.8073, + 8.5544, + 8.215, + 7.8348, + 7.4522, + 7.1063, + 6.8145, + 6.5775, + 6.3964, + 6.2819, + 6.2235, + 6.2045 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 1.8313, + 1.8325, + 1.839, + 1.8519, + 1.8723, + 1.9019, + 1.9406, + 1.9867, + 2.0379, + 2.0901, + 2.1385, + 2.1763, + 2.1985, + 2.2012, + 2.1845, + 2.1519, + 2.1072, + 2.0557, + 2.0022, + 1.9509, + 1.9054, + 1.8692, + 1.8441, + 1.83, + 1.8253 + ], + "Test Loss": [ + 1.8313, + 1.8325, + 1.839, + 1.8519, + 1.8723, + 1.9019, + 1.9406, + 1.9867, + 2.0379, + 2.0901, + 2.1385, + 2.1763, + 2.1985, + 2.2012, + 2.1845, + 2.1519, + 2.1072, + 2.0557, + 2.0022, + 1.9509, + 1.9054, + 1.8692, + 1.8441, + 1.83, + 1.8253 + ], + "Val PPL": [ + 6.2418, + 6.2496, + 6.2904, + 6.3718, + 6.5035, + 6.6984, + 6.9627, + 7.2913, + 7.6742, + 8.0857, + 8.4869, + 8.814, + 9.0114, + 9.0356, + 8.8861, + 8.6008, + 8.2253, + 7.8125, + 7.405, + 7.0349, + 6.722, + 6.4832, + 6.3224, + 6.2337, + 6.2046 + ], + "Test PPL": [ + 6.2418, + 6.2496, + 6.2904, + 6.3718, + 6.5035, + 6.6984, + 6.9627, + 7.2913, + 7.6742, + 8.0857, + 8.4869, + 8.814, + 9.0114, + 9.0356, + 8.8861, + 8.6008, + 8.2253, + 7.8125, + 7.405, + 7.0349, + 6.722, + 6.4832, + 6.3224, + 6.2337, + 6.2046 + ] + } +} \ No newline at end of file diff --git a/results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed20].json b/results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..90e11d808b4134ba61a6cebdd5e6523c10f4f77b --- /dev/null +++ b/results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.7797, + 3.7906, + 3.8317, + 3.9381, + 4.2034, + 4.7306, + 5.4334, + 6.119, + 6.6896, + 7.1002, + 7.3654, + 7.5113, + 7.5537, + 7.5061, + 7.3526, + 7.0649, + 6.637, + 6.0409, + 5.3436, + 4.6616, + 4.1733, + 3.9297, + 3.8288, + 3.789, + 3.7788 + ], + "Test Loss": [ + 3.7797, + 3.7906, + 3.8317, + 3.9381, + 4.2034, + 4.7306, + 5.4334, + 6.119, + 6.6896, + 7.1002, + 7.3654, + 7.5113, + 7.5537, + 7.5061, + 7.3526, + 7.0649, + 6.637, + 6.0409, + 5.3436, + 4.6616, + 4.1733, + 3.9297, + 3.8288, + 3.789, + 3.7788 + ], + "Val PPL": [ + 43.8044, + 44.2839, + 46.1409, + 51.3235, + 66.911, + 113.3587, + 228.9236, + 454.4198, + 804.0359, + 1212.1958, + 1580.3118, + 1828.5222, + 1907.803, + 1819.1002, + 1560.181, + 1170.1464, + 762.8036, + 420.2616, + 209.2674, + 105.802, + 64.9277, + 50.8893, + 46.007, + 44.214, + 43.762 + ], + "Test PPL": [ + 43.8044, + 44.2839, + 46.1409, + 51.3235, + 66.911, + 113.3587, + 228.9236, + 454.4198, + 804.0359, + 1212.1958, + 1580.3118, + 1828.5222, + 1907.803, + 1819.1002, + 1560.181, + 1170.1464, + 762.8036, + 420.2616, + 209.2674, + 105.802, + 64.9277, + 50.8893, + 46.007, + 44.214, + 43.762 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.7797, + 3.7786, + 3.778, + 3.7778, + 3.7776, + 3.7775, + 3.7777, + 3.7777, + 3.7779, + 3.7781, + 3.7781, + 3.7783, + 3.7783, + 3.7782, + 3.778, + 3.7778, + 3.7775, + 3.7771, + 3.7769, + 3.7767, + 3.7768, + 3.7768, + 3.7773, + 3.7779, + 3.7788 + ], + "Test Loss": [ + 3.7797, + 3.7786, + 3.778, + 3.7778, + 3.7776, + 3.7775, + 3.7777, + 3.7777, + 3.7779, + 3.7781, + 3.7781, + 3.7783, + 3.7783, + 3.7782, + 3.778, + 3.7778, + 3.7775, + 3.7771, + 3.7769, + 3.7767, + 3.7768, + 3.7768, + 3.7773, + 3.7779, + 3.7788 + ], + "Val PPL": [ + 43.8044, + 43.7545, + 43.7269, + 43.7178, + 43.7113, + 43.7066, + 43.714, + 43.717, + 43.724, + 43.7326, + 43.7344, + 43.7411, + 43.7395, + 43.7374, + 43.7273, + 43.7186, + 43.7052, + 43.691, + 43.6826, + 43.6735, + 43.6749, + 43.6764, + 43.696, + 43.7232, + 43.7615 + ], + "Test PPL": [ + 43.8044, + 43.7545, + 43.7269, + 43.7178, + 43.7113, + 43.7066, + 43.714, + 43.717, + 43.724, + 43.7326, + 43.7344, + 43.7411, + 43.7395, + 43.7374, + 43.7273, + 43.7186, + 43.7052, + 43.691, + 43.6826, + 43.6735, + 43.6749, + 43.6764, + 43.696, + 43.7232, + 43.7615 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.7797, + 3.7788, + 3.7782, + 3.7779, + 3.7776, + 3.7775, + 3.7778, + 3.7779, + 3.7779, + 3.7781, + 3.7784, + 3.7784, + 3.7784, + 3.7784, + 3.7784, + 3.7781, + 3.7777, + 3.7774, + 3.7769, + 3.7768, + 3.7768, + 3.7767, + 3.7771, + 3.7778, + 3.7788 + ], + "Test Loss": [ + 3.7797, + 3.7788, + 3.7782, + 3.7779, + 3.7776, + 3.7775, + 3.7778, + 3.7779, + 3.7779, + 3.7781, + 3.7784, + 3.7784, + 3.7784, + 3.7784, + 3.7784, + 3.7781, + 3.7777, + 3.7774, + 3.7769, + 3.7768, + 3.7768, + 3.7767, + 3.7771, + 3.7778, + 3.7788 + ], + "Val PPL": [ + 43.8044, + 43.762, + 43.7374, + 43.7222, + 43.7113, + 43.7079, + 43.7193, + 43.7224, + 43.725, + 43.7326, + 43.7463, + 43.745, + 43.7475, + 43.7466, + 43.7445, + 43.7307, + 43.714, + 43.7027, + 43.6822, + 43.6782, + 43.6742, + 43.6729, + 43.6897, + 43.7182, + 43.7626 + ], + "Test PPL": [ + 43.8044, + 43.762, + 43.7374, + 43.7222, + 43.7113, + 43.7079, + 43.7193, + 43.7224, + 43.725, + 43.7326, + 43.7463, + 43.745, + 43.7475, + 43.7466, + 43.7445, + 43.7307, + 43.714, + 43.7027, + 43.6822, + 43.6782, + 43.6742, + 43.6729, + 43.6897, + 43.7182, + 43.7626 + ] + } +} \ No newline at end of file diff --git a/results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed40].json b/results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..9246835dfc66dfdf78e988d090f45054cda4d434 --- /dev/null +++ b/results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.7797, + 3.7907, + 3.8303, + 3.9293, + 4.1696, + 4.644, + 5.3177, + 6.0439, + 6.6835, + 7.1573, + 7.4693, + 7.6402, + 7.6971, + 7.6378, + 7.4537, + 7.1299, + 6.6419, + 6.0031, + 5.2862, + 4.6149, + 4.1502, + 3.9225, + 3.8278, + 3.7898, + 3.7801 + ], + "Test Loss": [ + 3.7797, + 3.7907, + 3.8303, + 3.9293, + 4.1696, + 4.644, + 5.3177, + 6.0439, + 6.6835, + 7.1573, + 7.4693, + 7.6402, + 7.6971, + 7.6378, + 7.4537, + 7.1299, + 6.6419, + 6.0031, + 5.2862, + 4.6149, + 4.1502, + 3.9225, + 3.8278, + 3.7898, + 3.7801 + ], + "Val PPL": [ + 43.8044, + 44.2889, + 46.0777, + 50.8727, + 64.6869, + 103.9549, + 203.9185, + 421.5469, + 799.1372, + 1283.4189, + 1753.4587, + 2080.2185, + 2201.8494, + 2075.0886, + 1726.2944, + 1248.739, + 766.5537, + 404.6819, + 197.5862, + 100.9805, + 63.4488, + 50.528, + 45.9625, + 44.2469, + 43.8185 + ], + "Test PPL": [ + 43.8044, + 44.2889, + 46.0777, + 50.8727, + 64.6869, + 103.9549, + 203.9185, + 421.5469, + 799.1372, + 1283.4189, + 1753.4587, + 2080.2185, + 2201.8494, + 2075.0886, + 1726.2944, + 1248.739, + 766.5537, + 404.6819, + 197.5862, + 100.9805, + 63.4488, + 50.528, + 45.9625, + 44.2469, + 43.8185 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.7797, + 3.7787, + 3.778, + 3.7777, + 3.7776, + 3.7777, + 3.7776, + 3.778, + 3.7783, + 3.7785, + 3.7786, + 3.7787, + 3.7787, + 3.7786, + 3.7785, + 3.7783, + 3.7781, + 3.7778, + 3.7776, + 3.7775, + 3.7775, + 3.7776, + 3.7782, + 3.7789, + 3.7801 + ], + "Test Loss": [ + 3.7797, + 3.7787, + 3.778, + 3.7777, + 3.7776, + 3.7777, + 3.7776, + 3.778, + 3.7783, + 3.7785, + 3.7786, + 3.7787, + 3.7787, + 3.7786, + 3.7785, + 3.7783, + 3.7781, + 3.7778, + 3.7776, + 3.7775, + 3.7775, + 3.7776, + 3.7782, + 3.7789, + 3.7801 + ], + "Val PPL": [ + 43.8044, + 43.7582, + 43.7275, + 43.7137, + 43.7113, + 43.7148, + 43.7124, + 43.7304, + 43.7425, + 43.7482, + 43.7551, + 43.7603, + 43.7608, + 43.7532, + 43.749, + 43.7405, + 43.7309, + 43.7208, + 43.709, + 43.7086, + 43.7053, + 43.7107, + 43.7366, + 43.7697, + 43.8192 + ], + "Test PPL": [ + 43.8044, + 43.7582, + 43.7275, + 43.7137, + 43.7113, + 43.7148, + 43.7124, + 43.7304, + 43.7425, + 43.7482, + 43.7551, + 43.7603, + 43.7608, + 43.7532, + 43.749, + 43.7405, + 43.7309, + 43.7208, + 43.709, + 43.7086, + 43.7053, + 43.7107, + 43.7366, + 43.7697, + 43.8192 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.7797, + 3.7787, + 3.7782, + 3.7778, + 3.7777, + 3.7777, + 3.7778, + 3.778, + 3.7783, + 3.7784, + 3.7785, + 3.7787, + 3.7787, + 3.7786, + 3.7787, + 3.7784, + 3.7782, + 3.7777, + 3.7776, + 3.7774, + 3.7773, + 3.7775, + 3.7779, + 3.7787, + 3.78 + ], + "Test Loss": [ + 3.7797, + 3.7787, + 3.7782, + 3.7778, + 3.7777, + 3.7777, + 3.7778, + 3.778, + 3.7783, + 3.7784, + 3.7785, + 3.7787, + 3.7787, + 3.7786, + 3.7787, + 3.7784, + 3.7782, + 3.7777, + 3.7776, + 3.7774, + 3.7773, + 3.7775, + 3.7779, + 3.7787, + 3.78 + ], + "Val PPL": [ + 43.8044, + 43.7613, + 43.7357, + 43.7212, + 43.7145, + 43.7167, + 43.7206, + 43.7274, + 43.7396, + 43.7446, + 43.7508, + 43.7599, + 43.7583, + 43.7565, + 43.7603, + 43.7441, + 43.7354, + 43.7169, + 43.7104, + 43.701, + 43.6979, + 43.7067, + 43.7263, + 43.7603, + 43.8173 + ], + "Test PPL": [ + 43.8044, + 43.7613, + 43.7357, + 43.7212, + 43.7145, + 43.7167, + 43.7206, + 43.7274, + 43.7396, + 43.7446, + 43.7508, + 43.7599, + 43.7583, + 43.7565, + 43.7603, + 43.7441, + 43.7354, + 43.7169, + 43.7104, + 43.701, + 43.6979, + 43.7067, + 43.7263, + 43.7603, + 43.8173 + ] + } +} \ No newline at end of file diff --git a/results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed40].json b/results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..ecd80966295754da458f537ae766ac57fb430cf8 --- /dev/null +++ b/results/wt103/[finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.7788, + 3.7901, + 3.8303, + 3.9258, + 4.1305, + 4.4824, + 4.9383, + 5.393, + 5.7766, + 6.0798, + 6.2924, + 6.4169, + 6.4622, + 6.4267, + 6.3058, + 6.1005, + 5.8039, + 5.4153, + 4.9557, + 4.4933, + 4.1333, + 3.9271, + 3.8319, + 3.7913, + 3.7801 + ], + "Test Loss": [ + 3.7788, + 3.7901, + 3.8303, + 3.9258, + 4.1305, + 4.4824, + 4.9383, + 5.393, + 5.7766, + 6.0798, + 6.2924, + 6.4169, + 6.4622, + 6.4267, + 6.3058, + 6.1005, + 5.8039, + 5.4153, + 4.9557, + 4.4933, + 4.1333, + 3.9271, + 3.8319, + 3.7913, + 3.7801 + ], + "Val PPL": [ + 43.762, + 44.2619, + 46.0746, + 50.6927, + 62.2104, + 88.4455, + 139.5335, + 219.8683, + 322.6443, + 436.9423, + 540.4589, + 612.1013, + 640.4558, + 618.1393, + 547.7228, + 446.0865, + 331.5774, + 224.8112, + 141.9767, + 89.4132, + 62.3807, + 50.7595, + 46.1492, + 44.3119, + 43.8185 + ], + "Test PPL": [ + 43.762, + 44.2619, + 46.0746, + 50.6927, + 62.2104, + 88.4455, + 139.5335, + 219.8683, + 322.6443, + 436.9423, + 540.4589, + 612.1013, + 640.4558, + 618.1393, + 547.7228, + 446.0865, + 331.5774, + 224.8112, + 141.9767, + 89.4132, + 62.3807, + 50.7595, + 46.1492, + 44.3119, + 43.8185 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.7788, + 3.7779, + 3.7773, + 3.777, + 3.7768, + 3.7769, + 3.7769, + 3.7771, + 3.7774, + 3.7775, + 3.7777, + 3.7778, + 3.7777, + 3.7778, + 3.7777, + 3.7776, + 3.7775, + 3.7774, + 3.7773, + 3.7774, + 3.7774, + 3.7777, + 3.7782, + 3.7789, + 3.7801 + ], + "Test Loss": [ + 3.7788, + 3.7779, + 3.7773, + 3.777, + 3.7768, + 3.7769, + 3.7769, + 3.7771, + 3.7774, + 3.7775, + 3.7777, + 3.7778, + 3.7777, + 3.7778, + 3.7777, + 3.7776, + 3.7775, + 3.7774, + 3.7773, + 3.7774, + 3.7774, + 3.7777, + 3.7782, + 3.7789, + 3.7801 + ], + "Val PPL": [ + 43.762, + 43.7236, + 43.6964, + 43.684, + 43.6763, + 43.6801, + 43.6784, + 43.6911, + 43.7011, + 43.7074, + 43.7142, + 43.7208, + 43.7169, + 43.7177, + 43.7139, + 43.7089, + 43.7049, + 43.704, + 43.6989, + 43.7042, + 43.704, + 43.7168, + 43.7394, + 43.7667, + 43.8224 + ], + "Test PPL": [ + 43.762, + 43.7236, + 43.6964, + 43.684, + 43.6763, + 43.6801, + 43.6784, + 43.6911, + 43.7011, + 43.7074, + 43.7142, + 43.7208, + 43.7169, + 43.7177, + 43.7139, + 43.7089, + 43.7049, + 43.704, + 43.6989, + 43.7042, + 43.704, + 43.7168, + 43.7394, + 43.7667, + 43.8224 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.7788, + 3.778, + 3.7775, + 3.7772, + 3.7772, + 3.7771, + 3.7771, + 3.7772, + 3.7773, + 3.7775, + 3.7777, + 3.7778, + 3.7779, + 3.778, + 3.778, + 3.7779, + 3.7776, + 3.7776, + 3.7775, + 3.7773, + 3.7774, + 3.7775, + 3.778, + 3.7788, + 3.78 + ], + "Test Loss": [ + 3.7788, + 3.778, + 3.7775, + 3.7772, + 3.7772, + 3.7771, + 3.7771, + 3.7772, + 3.7773, + 3.7775, + 3.7777, + 3.7778, + 3.7779, + 3.778, + 3.778, + 3.7779, + 3.7776, + 3.7776, + 3.7775, + 3.7773, + 3.7774, + 3.7775, + 3.778, + 3.7788, + 3.78 + ], + "Val PPL": [ + 43.762, + 43.7304, + 43.7085, + 43.6953, + 43.6937, + 43.6879, + 43.6886, + 43.6934, + 43.6982, + 43.7083, + 43.7174, + 43.721, + 43.7243, + 43.7298, + 43.7276, + 43.7231, + 43.7128, + 43.709, + 43.7054, + 43.7, + 43.7042, + 43.7082, + 43.7297, + 43.7644, + 43.8182 + ], + "Test PPL": [ + 43.762, + 43.7304, + 43.7085, + 43.6953, + 43.6937, + 43.6879, + 43.6886, + 43.6934, + 43.6982, + 43.7083, + 43.7174, + 43.721, + 43.7243, + 43.7298, + 43.7276, + 43.7231, + 43.7128, + 43.709, + 43.7054, + 43.7, + 43.7042, + 43.7082, + 43.7297, + 43.7644, + 43.8182 + ] + } +} \ No newline at end of file diff --git a/results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed20].json b/results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..88b724ccbbab5a47899ca54a022f461f49c0b456 --- /dev/null +++ b/results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.7763, + 3.7849, + 3.819, + 3.9014, + 4.0911, + 4.4649, + 5.0366, + 5.7025, + 6.3205, + 6.8003, + 7.1166, + 7.2882, + 7.3269, + 7.238, + 7.0179, + 6.6497, + 6.1282, + 5.5114, + 4.8948, + 4.384, + 4.0532, + 3.8872, + 3.8118, + 3.7797, + 3.772 + ], + "Test Loss": [ + 3.7763, + 3.7849, + 3.819, + 3.9014, + 4.0911, + 4.4649, + 5.0366, + 5.7025, + 6.3205, + 6.8003, + 7.1166, + 7.2882, + 7.3269, + 7.238, + 7.0179, + 6.6497, + 6.1282, + 5.5114, + 4.8948, + 4.384, + 4.0532, + 3.8872, + 3.8118, + 3.7797, + 3.772 + ], + "Val PPL": [ + 43.6539, + 44.0312, + 45.557, + 49.4713, + 59.8057, + 86.9149, + 153.9471, + 299.6026, + 555.8394, + 898.116, + 1232.2578, + 1462.8993, + 1520.6183, + 1391.3263, + 1116.4198, + 772.5294, + 458.6231, + 247.5076, + 133.5945, + 80.1553, + 57.5824, + 48.7749, + 45.2299, + 43.804, + 43.4657 + ], + "Test PPL": [ + 43.6539, + 44.0312, + 45.557, + 49.4713, + 59.8057, + 86.9149, + 153.9471, + 299.6026, + 555.8394, + 898.116, + 1232.2578, + 1462.8993, + 1520.6183, + 1391.3263, + 1116.4198, + 772.5294, + 458.6231, + 247.5076, + 133.5945, + 80.1553, + 57.5824, + 48.7749, + 45.2299, + 43.804, + 43.4657 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.7763, + 3.7749, + 3.775, + 3.7764, + 3.7794, + 3.784, + 3.7901, + 3.7973, + 3.8052, + 3.8131, + 3.8203, + 3.8247, + 3.8252, + 3.8221, + 3.8151, + 3.8062, + 3.7975, + 3.79, + 3.7831, + 3.7784, + 3.7746, + 3.7721, + 3.7707, + 3.7707, + 3.772 + ], + "Test Loss": [ + 3.7763, + 3.7749, + 3.775, + 3.7764, + 3.7794, + 3.784, + 3.7901, + 3.7973, + 3.8052, + 3.8131, + 3.8203, + 3.8247, + 3.8252, + 3.8221, + 3.8151, + 3.8062, + 3.7975, + 3.79, + 3.7831, + 3.7784, + 3.7746, + 3.7721, + 3.7707, + 3.7707, + 3.772 + ], + "Val PPL": [ + 43.6539, + 43.5914, + 43.5954, + 43.6599, + 43.7912, + 43.9902, + 44.2597, + 44.5815, + 44.9345, + 45.2891, + 45.62, + 45.8181, + 45.8427, + 45.7012, + 45.3805, + 44.9803, + 44.5899, + 44.2569, + 43.9527, + 43.748, + 43.5782, + 43.4707, + 43.4086, + 43.4106, + 43.4655 + ], + "Test PPL": [ + 43.6539, + 43.5914, + 43.5954, + 43.6599, + 43.7912, + 43.9902, + 44.2597, + 44.5815, + 44.9345, + 45.2891, + 45.62, + 45.8181, + 45.8427, + 45.7012, + 45.3805, + 44.9803, + 44.5899, + 44.2569, + 43.9527, + 43.748, + 43.5782, + 43.4707, + 43.4086, + 43.4106, + 43.4655 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.7763, + 3.775, + 3.7752, + 3.7765, + 3.7791, + 3.7829, + 3.7885, + 3.7953, + 3.8033, + 3.8123, + 3.8218, + 3.8296, + 3.834, + 3.8323, + 3.8258, + 3.8159, + 3.805, + 3.7952, + 3.787, + 3.7806, + 3.7758, + 3.7726, + 3.7708, + 3.7706, + 3.772 + ], + "Test Loss": [ + 3.7763, + 3.775, + 3.7752, + 3.7765, + 3.7791, + 3.7829, + 3.7885, + 3.7953, + 3.8033, + 3.8123, + 3.8218, + 3.8296, + 3.834, + 3.8323, + 3.8258, + 3.8159, + 3.805, + 3.7952, + 3.787, + 3.7806, + 3.7758, + 3.7726, + 3.7708, + 3.7706, + 3.772 + ], + "Val PPL": [ + 43.6539, + 43.5981, + 43.6045, + 43.6633, + 43.7757, + 43.9426, + 44.1917, + 44.4909, + 44.8492, + 45.2524, + 45.6876, + 46.0446, + 46.2454, + 46.1704, + 45.868, + 45.4154, + 44.9268, + 44.4878, + 44.1235, + 43.843, + 43.6327, + 43.4934, + 43.4146, + 43.4051, + 43.4672 + ], + "Test PPL": [ + 43.6539, + 43.5981, + 43.6045, + 43.6633, + 43.7757, + 43.9426, + 44.1917, + 44.4909, + 44.8492, + 45.2524, + 45.6876, + 46.0446, + 46.2454, + 46.1704, + 45.868, + 45.4154, + 44.9268, + 44.4878, + 44.1235, + 43.843, + 43.6327, + 43.4934, + 43.4146, + 43.4051, + 43.4672 + ] + } +} \ No newline at end of file diff --git a/results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed40].json b/results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..182b5af997b37d51550e72601a7972d051e46f77 --- /dev/null +++ b/results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.7763, + 3.7864, + 3.8215, + 3.9034, + 4.0981, + 4.5156, + 5.1408, + 5.8187, + 6.4242, + 6.891, + 7.193, + 7.3601, + 7.415, + 7.3632, + 7.1906, + 6.8633, + 6.3408, + 5.6586, + 4.9489, + 4.387, + 4.0501, + 3.8895, + 3.8157, + 3.7831, + 3.7744 + ], + "Test Loss": [ + 3.7763, + 3.7864, + 3.8215, + 3.9034, + 4.0981, + 4.5156, + 5.1408, + 5.8187, + 6.4242, + 6.891, + 7.193, + 7.3601, + 7.415, + 7.3632, + 7.1906, + 6.8633, + 6.3408, + 5.6586, + 4.9489, + 4.387, + 4.0501, + 3.8895, + 3.8157, + 3.7831, + 3.7744 + ], + "Val PPL": [ + 43.6539, + 44.0955, + 45.6721, + 49.5684, + 60.2275, + 91.4348, + 170.8556, + 336.5246, + 616.5915, + 983.3996, + 1330.1317, + 1571.9784, + 1660.6323, + 1576.929, + 1326.9382, + 956.4936, + 567.2298, + 286.7383, + 141.0254, + 80.3963, + 57.4059, + 48.887, + 45.4101, + 43.9506, + 43.5705 + ], + "Test PPL": [ + 43.6539, + 44.0955, + 45.6721, + 49.5684, + 60.2275, + 91.4348, + 170.8556, + 336.5246, + 616.5915, + 983.3996, + 1330.1317, + 1571.9784, + 1660.6323, + 1576.929, + 1326.9382, + 956.4936, + 567.2298, + 286.7383, + 141.0254, + 80.3963, + 57.4059, + 48.887, + 45.4101, + 43.9506, + 43.5705 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.7763, + 3.7752, + 3.7748, + 3.7755, + 3.7765, + 3.7781, + 3.7799, + 3.782, + 3.7841, + 3.786, + 3.7875, + 3.7883, + 3.7885, + 3.788, + 3.7867, + 3.7849, + 3.7827, + 3.7804, + 3.7782, + 3.7762, + 3.7745, + 3.7733, + 3.7728, + 3.773, + 3.7743 + ], + "Test Loss": [ + 3.7763, + 3.7752, + 3.7748, + 3.7755, + 3.7765, + 3.7781, + 3.7799, + 3.782, + 3.7841, + 3.786, + 3.7875, + 3.7883, + 3.7885, + 3.788, + 3.7867, + 3.7849, + 3.7827, + 3.7804, + 3.7782, + 3.7762, + 3.7745, + 3.7733, + 3.7728, + 3.773, + 3.7743 + ], + "Val PPL": [ + 43.6539, + 43.6045, + 43.5888, + 43.6177, + 43.6621, + 43.7339, + 43.8127, + 43.9037, + 43.9948, + 44.078, + 44.1469, + 44.1824, + 44.1904, + 44.1685, + 44.1127, + 44.0294, + 43.9346, + 43.8332, + 43.7351, + 43.648, + 43.5741, + 43.5248, + 43.5037, + 43.511, + 43.5668 + ], + "Test PPL": [ + 43.6539, + 43.6045, + 43.5888, + 43.6177, + 43.6621, + 43.7339, + 43.8127, + 43.9037, + 43.9948, + 44.078, + 44.1469, + 44.1824, + 44.1904, + 44.1685, + 44.1127, + 44.0294, + 43.9346, + 43.8332, + 43.7351, + 43.648, + 43.5741, + 43.5248, + 43.5037, + 43.511, + 43.5668 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.7763, + 3.7752, + 3.7751, + 3.7757, + 3.7767, + 3.7784, + 3.7802, + 3.7821, + 3.7843, + 3.7864, + 3.7882, + 3.7895, + 3.7902, + 3.7898, + 3.7888, + 3.7869, + 3.7844, + 3.7819, + 3.7794, + 3.777, + 3.775, + 3.7735, + 3.7729, + 3.773, + 3.7744 + ], + "Test Loss": [ + 3.7763, + 3.7752, + 3.7751, + 3.7757, + 3.7767, + 3.7784, + 3.7802, + 3.7821, + 3.7843, + 3.7864, + 3.7882, + 3.7895, + 3.7902, + 3.7898, + 3.7888, + 3.7869, + 3.7844, + 3.7819, + 3.7794, + 3.777, + 3.775, + 3.7735, + 3.7729, + 3.773, + 3.7744 + ], + "Val PPL": [ + 43.6539, + 43.6078, + 43.604, + 43.6273, + 43.6736, + 43.7464, + 43.8255, + 43.9097, + 44.0061, + 44.0962, + 44.178, + 44.2329, + 44.2648, + 44.2487, + 44.2013, + 44.1173, + 44.0093, + 43.9001, + 43.7883, + 43.6849, + 43.5961, + 43.5338, + 43.5051, + 43.5089, + 43.5714 + ], + "Test PPL": [ + 43.6539, + 43.6078, + 43.604, + 43.6273, + 43.6736, + 43.7464, + 43.8255, + 43.9097, + 44.0061, + 44.0962, + 44.178, + 44.2329, + 44.2648, + 44.2487, + 44.2013, + 44.1173, + 44.0093, + 43.9001, + 43.7883, + 43.6849, + 43.5961, + 43.5338, + 43.5051, + 43.5089, + 43.5714 + ] + } +} \ No newline at end of file diff --git a/results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed40].json b/results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..231859c3d4b130a439c1b23f18e38f5d141f1154 --- /dev/null +++ b/results/wt103/[finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.772, + 3.781, + 3.8146, + 3.8879, + 4.0436, + 4.3519, + 4.8245, + 5.3805, + 5.9088, + 6.3511, + 6.6639, + 6.8473, + 6.9042, + 6.8418, + 6.6558, + 6.3434, + 5.9024, + 5.3683, + 4.8107, + 4.3374, + 4.0397, + 3.8873, + 3.8157, + 3.7828, + 3.7744 + ], + "Test Loss": [ + 3.772, + 3.781, + 3.8146, + 3.8879, + 4.0436, + 4.3519, + 4.8245, + 5.3805, + 5.9088, + 6.3511, + 6.6639, + 6.8473, + 6.9042, + 6.8418, + 6.6558, + 6.3434, + 5.9024, + 5.3683, + 4.8107, + 4.3374, + 4.0397, + 3.8873, + 3.8157, + 3.7828, + 3.7744 + ], + "Val PPL": [ + 43.4657, + 43.8606, + 45.357, + 48.8089, + 57.0303, + 77.6282, + 124.5271, + 217.1227, + 368.2546, + 573.1353, + 783.6328, + 941.3076, + 996.5, + 936.1636, + 777.2601, + 568.7109, + 365.9022, + 214.5064, + 122.8224, + 76.5054, + 56.8111, + 48.7794, + 45.4074, + 43.9379, + 43.5705 + ], + "Test PPL": [ + 43.4657, + 43.8606, + 45.357, + 48.8089, + 57.0303, + 77.6282, + 124.5271, + 217.1227, + 368.2546, + 573.1353, + 783.6328, + 941.3076, + 996.5, + 936.1636, + 777.2601, + 568.7109, + 365.9022, + 214.5064, + 122.8224, + 76.5054, + 56.8111, + 48.7794, + 45.4074, + 43.9379, + 43.5705 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.772, + 3.7711, + 3.7709, + 3.7715, + 3.7728, + 3.7744, + 3.7762, + 3.7785, + 3.7807, + 3.7829, + 3.785, + 3.7868, + 3.7879, + 3.7879, + 3.7867, + 3.7846, + 3.7823, + 3.7798, + 3.7776, + 3.7757, + 3.7742, + 3.7732, + 3.7728, + 3.7732, + 3.7744 + ], + "Test Loss": [ + 3.772, + 3.7711, + 3.7709, + 3.7715, + 3.7728, + 3.7744, + 3.7762, + 3.7785, + 3.7807, + 3.7829, + 3.785, + 3.7868, + 3.7879, + 3.7879, + 3.7867, + 3.7846, + 3.7823, + 3.7798, + 3.7776, + 3.7757, + 3.7742, + 3.7732, + 3.7728, + 3.7732, + 3.7744 + ], + "Val PPL": [ + 43.4657, + 43.4265, + 43.421, + 43.4464, + 43.4997, + 43.5701, + 43.6514, + 43.7484, + 43.8477, + 43.9448, + 44.0372, + 44.1168, + 44.1619, + 44.1637, + 44.1093, + 44.0177, + 43.915, + 43.8062, + 43.709, + 43.6282, + 43.5605, + 43.5198, + 43.501, + 43.517, + 43.571 + ], + "Test PPL": [ + 43.4657, + 43.4265, + 43.421, + 43.4464, + 43.4997, + 43.5701, + 43.6514, + 43.7484, + 43.8477, + 43.9448, + 44.0372, + 44.1168, + 44.1619, + 44.1637, + 44.1093, + 44.0177, + 43.915, + 43.8062, + 43.709, + 43.6282, + 43.5605, + 43.5198, + 43.501, + 43.517, + 43.571 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.772, + 3.7712, + 3.7714, + 3.7721, + 3.7734, + 3.7749, + 3.7768, + 3.7791, + 3.7816, + 3.7839, + 3.7863, + 3.7884, + 3.7901, + 3.7909, + 3.7906, + 3.789, + 3.7864, + 3.7834, + 3.7804, + 3.7776, + 3.7753, + 3.7739, + 3.773, + 3.7732, + 3.7744 + ], + "Test Loss": [ + 3.772, + 3.7712, + 3.7714, + 3.7721, + 3.7734, + 3.7749, + 3.7768, + 3.7791, + 3.7816, + 3.7839, + 3.7863, + 3.7884, + 3.7901, + 3.7909, + 3.7906, + 3.789, + 3.7864, + 3.7834, + 3.7804, + 3.7776, + 3.7753, + 3.7739, + 3.773, + 3.7732, + 3.7744 + ], + "Val PPL": [ + 43.4657, + 43.4334, + 43.4417, + 43.4718, + 43.5259, + 43.5949, + 43.6767, + 43.7752, + 43.8855, + 43.9892, + 44.0908, + 44.1838, + 44.2587, + 44.2979, + 44.2814, + 44.2115, + 44.0973, + 43.9637, + 43.8337, + 43.7105, + 43.6122, + 43.5487, + 43.5125, + 43.5184, + 43.5699 + ], + "Test PPL": [ + 43.4657, + 43.4334, + 43.4417, + 43.4718, + 43.5259, + 43.5949, + 43.6767, + 43.7752, + 43.8855, + 43.9892, + 44.0908, + 44.1838, + 44.2587, + 44.2979, + 44.2814, + 44.2115, + 44.0973, + 43.9637, + 43.8337, + 43.7105, + 43.6122, + 43.5487, + 43.5125, + 43.5184, + 43.5699 + ] + } +} \ No newline at end of file diff --git a/results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json b/results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json new file mode 100644 index 0000000000000000000000000000000000000000..91fa22d244c9d0a8ec89735776b7fff325345045 --- /dev/null +++ b/results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.7708, + 3.7766, + 3.8032, + 3.8632, + 3.996, + 4.2821, + 4.7574, + 5.3523, + 5.9645, + 6.4916, + 6.8723, + 7.0968, + 7.1675, + 7.0859, + 6.8584, + 6.4815, + 5.9688, + 5.3762, + 4.778, + 4.2875, + 3.9974, + 3.8631, + 3.8036, + 3.778, + 3.7732 + ], + "Test Loss": [ + 3.7708, + 3.7766, + 3.8032, + 3.8632, + 3.996, + 4.2821, + 4.7574, + 5.3523, + 5.9645, + 6.4916, + 6.8723, + 7.0968, + 7.1675, + 7.0859, + 6.8584, + 6.4815, + 5.9688, + 5.3762, + 4.778, + 4.2875, + 3.9974, + 3.8631, + 3.8036, + 3.778, + 3.7732 + ], + "Val PPL": [ + 43.4158, + 43.668, + 44.8434, + 47.6196, + 54.3809, + 72.3896, + 116.4375, + 211.1005, + 389.3553, + 659.5502, + 965.1559, + 1208.1375, + 1296.5614, + 1194.9445, + 951.8309, + 652.9417, + 391.0258, + 216.2047, + 118.8686, + 72.7851, + 54.4551, + 47.6128, + 44.8614, + 43.7299, + 43.5205 + ], + "Test PPL": [ + 43.4158, + 43.668, + 44.8434, + 47.6196, + 54.3809, + 72.3896, + 116.4375, + 211.1005, + 389.3553, + 659.5502, + 965.1559, + 1208.1375, + 1296.5614, + 1194.9445, + 951.8309, + 652.9417, + 391.0258, + 216.2047, + 118.8686, + 72.7851, + 54.4551, + 47.6128, + 44.8614, + 43.7299, + 43.5205 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.7708, + 3.7705, + 3.7723, + 3.7759, + 3.7808, + 3.7872, + 3.7946, + 3.8032, + 3.8129, + 3.8231, + 3.8325, + 3.8388, + 3.8412, + 3.8392, + 3.8327, + 3.8233, + 3.8124, + 3.8018, + 3.793, + 3.7861, + 3.7806, + 3.7762, + 3.7733, + 3.7723, + 3.7733 + ], + "Test Loss": [ + 3.7708, + 3.7705, + 3.7723, + 3.7759, + 3.7808, + 3.7872, + 3.7946, + 3.8032, + 3.8129, + 3.8231, + 3.8325, + 3.8388, + 3.8412, + 3.8392, + 3.8327, + 3.8233, + 3.8124, + 3.8018, + 3.793, + 3.7861, + 3.7806, + 3.7762, + 3.7733, + 3.7723, + 3.7733 + ], + "Val PPL": [ + 43.4158, + 43.4, + 43.4804, + 43.6357, + 43.853, + 44.1312, + 44.4605, + 44.8437, + 45.2815, + 45.7478, + 46.1757, + 46.4682, + 46.5809, + 46.4874, + 46.1878, + 45.7555, + 45.258, + 44.7839, + 44.3878, + 44.0831, + 43.8431, + 43.648, + 43.5256, + 43.4786, + 43.5225 + ], + "Test PPL": [ + 43.4158, + 43.4, + 43.4804, + 43.6357, + 43.853, + 44.1312, + 44.4605, + 44.8437, + 45.2815, + 45.7478, + 46.1757, + 46.4682, + 46.5809, + 46.4874, + 46.1878, + 45.7555, + 45.258, + 44.7839, + 44.3878, + 44.0831, + 43.8431, + 43.648, + 43.5256, + 43.4786, + 43.5225 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.7708, + 3.7708, + 3.7727, + 3.7761, + 3.7809, + 3.7869, + 3.7941, + 3.8021, + 3.811, + 3.8209, + 3.8306, + 3.8383, + 3.8426, + 3.8429, + 3.8384, + 3.8299, + 3.8187, + 3.8073, + 3.7973, + 3.7891, + 3.7824, + 3.7771, + 3.7736, + 3.7722, + 3.7733 + ], + "Test Loss": [ + 3.7708, + 3.7708, + 3.7727, + 3.7761, + 3.7809, + 3.7869, + 3.7941, + 3.8021, + 3.811, + 3.8209, + 3.8306, + 3.8383, + 3.8426, + 3.8429, + 3.8384, + 3.8299, + 3.8187, + 3.8073, + 3.7973, + 3.7891, + 3.7824, + 3.7771, + 3.7736, + 3.7722, + 3.7733 + ], + "Val PPL": [ + 43.4158, + 43.4149, + 43.4977, + 43.6451, + 43.8538, + 44.1197, + 44.4374, + 44.7946, + 45.1936, + 45.646, + 46.089, + 46.4466, + 46.6461, + 46.6618, + 46.4516, + 46.0576, + 45.5462, + 45.0273, + 44.5818, + 44.2161, + 43.9226, + 43.691, + 43.5385, + 43.4771, + 43.5226 + ], + "Test PPL": [ + 43.4158, + 43.4149, + 43.4977, + 43.6451, + 43.8538, + 44.1197, + 44.4374, + 44.7946, + 45.1936, + 45.646, + 46.089, + 46.4466, + 46.6461, + 46.6618, + 46.4516, + 46.0576, + 45.5462, + 45.0273, + 44.5818, + 44.2161, + 43.9226, + 43.691, + 43.5385, + 43.4771, + 43.5226 + ] + } +} \ No newline at end of file diff --git a/results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json b/results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..b1997d398ae21741d8b1c323b078258e316707c3 --- /dev/null +++ b/results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.7708, + 3.7775, + 3.8035, + 3.8578, + 3.9633, + 4.1809, + 4.6052, + 5.2373, + 5.909, + 6.4853, + 6.8959, + 7.1214, + 7.1909, + 7.1042, + 6.8689, + 6.4732, + 5.9211, + 5.2717, + 4.6474, + 4.199, + 3.9648, + 3.8562, + 3.8024, + 3.7777, + 3.7726 + ], + "Test Loss": [ + 3.7708, + 3.7775, + 3.8035, + 3.8578, + 3.9633, + 4.1809, + 4.6052, + 5.2373, + 5.909, + 6.4853, + 6.8959, + 7.1214, + 7.1909, + 7.1042, + 6.8689, + 6.4732, + 5.9211, + 5.2717, + 4.6474, + 4.199, + 3.9648, + 3.8562, + 3.8024, + 3.7777, + 3.7726 + ], + "Val PPL": [ + 43.4158, + 43.7052, + 44.8575, + 47.3613, + 52.629, + 65.4218, + 100.0023, + 188.1658, + 368.3417, + 655.4619, + 988.2178, + 1238.215, + 1327.347, + 1217.1289, + 961.8572, + 647.5296, + 372.8371, + 194.7552, + 104.3143, + 66.6186, + 52.7093, + 47.2841, + 44.8102, + 43.715, + 43.4947 + ], + "Test PPL": [ + 43.4158, + 43.7052, + 44.8575, + 47.3613, + 52.629, + 65.4218, + 100.0023, + 188.1658, + 368.3417, + 655.4619, + 988.2178, + 1238.215, + 1327.347, + 1217.1289, + 961.8572, + 647.5296, + 372.8371, + 194.7552, + 104.3143, + 66.6186, + 52.7093, + 47.2841, + 44.8102, + 43.715, + 43.4947 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.7708, + 3.7705, + 3.7723, + 3.7758, + 3.781, + 3.7872, + 3.7949, + 3.8037, + 3.8134, + 3.8235, + 3.8327, + 3.8385, + 3.8397, + 3.8358, + 3.8285, + 3.819, + 3.8095, + 3.8005, + 3.7927, + 3.7856, + 3.7797, + 3.7753, + 3.7724, + 3.7715, + 3.7726 + ], + "Test Loss": [ + 3.7708, + 3.7705, + 3.7723, + 3.7758, + 3.781, + 3.7872, + 3.7949, + 3.8037, + 3.8134, + 3.8235, + 3.8327, + 3.8385, + 3.8397, + 3.8358, + 3.8285, + 3.819, + 3.8095, + 3.8005, + 3.7927, + 3.7856, + 3.7797, + 3.7753, + 3.7724, + 3.7715, + 3.7726 + ], + "Val PPL": [ + 43.4158, + 43.4015, + 43.4788, + 43.6311, + 43.8583, + 44.1326, + 44.4755, + 44.8667, + 45.3039, + 45.7641, + 46.1851, + 46.458, + 46.5105, + 46.3323, + 45.9931, + 45.5609, + 45.1269, + 44.7242, + 44.3739, + 44.0608, + 43.8032, + 43.6097, + 43.4827, + 43.4444, + 43.4914 + ], + "Test PPL": [ + 43.4158, + 43.4015, + 43.4788, + 43.6311, + 43.8583, + 44.1326, + 44.4755, + 44.8667, + 45.3039, + 45.7641, + 46.1851, + 46.458, + 46.5105, + 46.3323, + 45.9931, + 45.5609, + 45.1269, + 44.7242, + 44.3739, + 44.0608, + 43.8032, + 43.6097, + 43.4827, + 43.4444, + 43.4914 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.7708, + 3.7705, + 3.772, + 3.7752, + 3.7798, + 3.7857, + 3.7926, + 3.8003, + 3.8093, + 3.8192, + 3.8289, + 3.8366, + 3.8401, + 3.8388, + 3.8326, + 3.8234, + 3.8133, + 3.8034, + 3.7946, + 3.7869, + 3.7802, + 3.7755, + 3.7723, + 3.7712, + 3.7726 + ], + "Test Loss": [ + 3.7708, + 3.7705, + 3.772, + 3.7752, + 3.7798, + 3.7857, + 3.7926, + 3.8003, + 3.8093, + 3.8192, + 3.8289, + 3.8366, + 3.8401, + 3.8388, + 3.8326, + 3.8234, + 3.8133, + 3.8034, + 3.7946, + 3.7869, + 3.7802, + 3.7755, + 3.7723, + 3.7712, + 3.7726 + ], + "Val PPL": [ + 43.4158, + 43.4021, + 43.4649, + 43.6041, + 43.8094, + 44.0649, + 44.3704, + 44.7168, + 45.1187, + 45.5677, + 46.0124, + 46.3664, + 46.5299, + 46.4699, + 46.1837, + 45.7611, + 45.2983, + 44.8546, + 44.4613, + 44.1204, + 43.8252, + 43.6189, + 43.4782, + 43.4312, + 43.4908 + ], + "Test PPL": [ + 43.4158, + 43.4021, + 43.4649, + 43.6041, + 43.8094, + 44.0649, + 44.3704, + 44.7168, + 45.1187, + 45.5677, + 46.0124, + 46.3664, + 46.5299, + 46.4699, + 46.1837, + 45.7611, + 45.2983, + 44.8546, + 44.4613, + 44.1204, + 43.8252, + 43.6189, + 43.4782, + 43.4312, + 43.4908 + ] + } +} \ No newline at end of file diff --git a/results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json b/results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json new file mode 100644 index 0000000000000000000000000000000000000000..23c627784097a80a65d40c2e13eef0a4ca983fd2 --- /dev/null +++ b/results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json @@ -0,0 +1,332 @@ +{ + "Naive": { + "Val Loss": [ + 3.7732, + 3.7782, + 3.8037, + 3.8613, + 3.9809, + 4.2188, + 4.6183, + 5.1343, + 5.6801, + 6.1707, + 6.5484, + 6.7732, + 6.8488, + 6.7817, + 6.5694, + 6.2093, + 5.7204, + 5.1587, + 4.6183, + 4.2148, + 3.978, + 3.86, + 3.8028, + 3.7773, + 3.7726 + ], + "Test Loss": [ + 3.7732, + 3.7782, + 3.8037, + 3.8613, + 3.9809, + 4.2188, + 4.6183, + 5.1343, + 5.6801, + 6.1707, + 6.5484, + 6.7732, + 6.8488, + 6.7817, + 6.5694, + 6.2093, + 5.7204, + 5.1587, + 4.6183, + 4.2148, + 3.978, + 3.86, + 3.8028, + 3.7773, + 3.7726 + ], + "Val PPL": [ + 43.5205, + 43.7389, + 44.8684, + 47.5254, + 53.5631, + 67.9516, + 101.3244, + 169.7525, + 292.9663, + 478.5013, + 698.1489, + 874.0892, + 942.7917, + 881.5343, + 712.9306, + 497.3352, + 305.0206, + 173.9318, + 101.3179, + 67.6791, + 53.4082, + 47.4639, + 44.8261, + 43.6994, + 43.4947 + ], + "Test PPL": [ + 43.5205, + 43.7389, + 44.8684, + 47.5254, + 53.5631, + 67.9516, + 101.3244, + 169.7525, + 292.9663, + 478.5013, + 698.1489, + 874.0892, + 942.7917, + 881.5343, + 712.9306, + 497.3352, + 305.0206, + 173.9318, + 101.3179, + 67.6791, + 53.4082, + 47.4639, + 44.8261, + 43.6994, + 43.4947 + ] + }, + "permu_head_init_ortho_no_opt": { + "Val Loss": [ + 3.7732, + 3.7724, + 3.7731, + 3.775, + 3.7781, + 3.7824, + 3.7876, + 3.7943, + 3.8035, + 3.815, + 3.8261, + 3.8343, + 3.8387, + 3.8376, + 3.8324, + 3.8227, + 3.8108, + 3.7984, + 3.7892, + 3.7827, + 3.7779, + 3.7745, + 3.7722, + 3.7715, + 3.7726 + ], + "Test Loss": [ + 3.7732, + 3.7724, + 3.7731, + 3.775, + 3.7781, + 3.7824, + 3.7876, + 3.7943, + 3.8035, + 3.815, + 3.8261, + 3.8343, + 3.8387, + 3.8376, + 3.8324, + 3.8227, + 3.8108, + 3.7984, + 3.7892, + 3.7827, + 3.7779, + 3.7745, + 3.7722, + 3.7715, + 3.7726 + ], + "Val PPL": [ + 43.5205, + 43.4852, + 43.5162, + 43.5961, + 43.734, + 43.9195, + 44.1503, + 44.4454, + 44.8584, + 45.3765, + 45.8831, + 46.262, + 46.4654, + 46.415, + 46.1754, + 45.7261, + 45.1854, + 44.6309, + 44.2205, + 43.9335, + 43.7262, + 43.5748, + 43.4772, + 43.4467, + 43.4922 + ], + "Test PPL": [ + 43.5205, + 43.4852, + 43.5162, + 43.5961, + 43.734, + 43.9195, + 44.1503, + 44.4454, + 44.8584, + 45.3765, + 45.8831, + 46.262, + 46.4654, + 46.415, + 46.1754, + 45.7261, + 45.1854, + 44.6309, + 44.2205, + 43.9335, + 43.7262, + 43.5748, + 43.4772, + 43.4467, + 43.4922 + ] + }, + "permu_head_init_ortho_opt": { + "Val Loss": [ + 3.7732, + 3.7727, + 3.7735, + 3.7757, + 3.779, + 3.7832, + 3.7884, + 3.7949, + 3.8037, + 3.8151, + 3.8275, + 3.8375, + 3.8444, + 3.846, + 3.8423, + 3.8338, + 3.8211, + 3.8066, + 3.7944, + 3.7857, + 3.7799, + 3.7756, + 3.7726, + 3.7717, + 3.7725 + ], + "Test Loss": [ + 3.7732, + 3.7727, + 3.7735, + 3.7757, + 3.779, + 3.7832, + 3.7884, + 3.7949, + 3.8037, + 3.8151, + 3.8275, + 3.8375, + 3.8444, + 3.846, + 3.8423, + 3.8338, + 3.8211, + 3.8066, + 3.7944, + 3.7857, + 3.7799, + 3.7756, + 3.7726, + 3.7717, + 3.7725 + ], + "Val PPL": [ + 43.5205, + 43.4977, + 43.5333, + 43.6285, + 43.7709, + 43.9554, + 44.1856, + 44.4757, + 44.8675, + 45.3793, + 45.9496, + 46.4081, + 46.7293, + 46.8053, + 46.6339, + 46.2369, + 45.6554, + 44.9978, + 44.4501, + 44.0681, + 43.8112, + 43.6238, + 43.4927, + 43.454, + 43.4906 + ], + "Test PPL": [ + 43.5205, + 43.4977, + 43.5333, + 43.6285, + 43.7709, + 43.9554, + 44.1856, + 44.4757, + 44.8675, + 45.3793, + 45.9496, + 46.4081, + 46.7293, + 46.8053, + 46.6339, + 46.2369, + 45.6554, + 44.9978, + 44.4501, + 44.0681, + 43.8112, + 43.6238, + 43.4927, + 43.454, + 43.4906 + ] + } +} \ No newline at end of file diff --git a/scripts/azure.sh b/scripts/azure.sh new file mode 100644 index 0000000000000000000000000000000000000000..7f8b7a8338d94d72008930d391d8f134b48a5313 --- /dev/null +++ b/scripts/azure.sh @@ -0,0 +1,8 @@ +git add README.md +git add plots +git add results +git add scripts +git add src +git add temp +git add pyproject.toml +git add requirements.txt diff --git a/scripts/cifar100/attn_match.sh b/scripts/cifar100/attn_match.sh new file mode 100644 index 0000000000000000000000000000000000000000..12ba349b1a20e4515ac918d1092820f635b49cb6 --- /dev/null +++ b/scripts/cifar100/attn_match.sh @@ -0,0 +1,9 @@ +CUDA_VISIBLE_DEVICES=0 python src/imagenet/att_matching.py --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 + +CUDA_VISIBLE_DEVICES=0 python src/imagenet/att_matching.py --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-seed0/best_16380 \ + --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 \ No newline at end of file diff --git a/scripts/cifar100/finetune.sh b/scripts/cifar100/finetune.sh new file mode 100644 index 0000000000000000000000000000000000000000..c2f296d6303e473734c2ba5db31761c6906f39e6 --- /dev/null +++ b/scripts/cifar100/finetune.sh @@ -0,0 +1,17 @@ +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/imagenet/finetune.py --input-size 32 --data-set CIFAR10\ + --model-path /mnt/d/NCKH/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 /mnt/d/NCKH/datasets/cifar10 --save-dir /mnt/d/NCKH/weights/lmc/cifar10/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/imagenet/finetune.py --input-size 32 --data-set CIFAR10\ + --model-path /mnt/d/NCKH/weights/lmc/cifar10/lr0.005-learnable-epochs50-batch128-seed0/best_15990 \ + --lr 5e-3 --epochs 50 --batch-size 128 --seed 20 --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 /mnt/d/NCKH/datasets/cifar10 --save-dir /mnt/d/NCKH/weights/lmc/cifar10/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/imagenet/finetune.py --input-size 32 --data-set CIFAR10\ + --model-path /mnt/d/NCKH/weights/lmc/cifar10/lr0.005-learnable-epochs50-batch128-seed0/best_15990 \ + --lr 5e-3 --epochs 50 --batch-size 128 --seed 40 --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 /mnt/d/NCKH/datasets/cifar10 --save-dir /mnt/d/NCKH/weights/lmc/cifar10/vit-finetune diff --git a/scripts/cifar100/train_model.sh b/scripts/cifar100/train_model.sh new file mode 100644 index 0000000000000000000000000000000000000000..e004063954074f7ae00a3df140343ac746705280 --- /dev/null +++ b/scripts/cifar100/train_model.sh @@ -0,0 +1,13 @@ +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 "learnable" --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 /mnt/d/NCKH/weights/lmc/cifar10 --data-path /mnt/d/NCKH/datasets/cifar10 + +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 "learnable" --num-labels 10\ + --lr 5e-4 --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 /mnt/d/NCKH/weights/lmc/cifar10 --data-path /mnt/d/NCKH/datasets/cifar10 \ No newline at end of file diff --git a/scripts/enwik8/att_matching.sh b/scripts/enwik8/att_matching.sh new file mode 100644 index 0000000000000000000000000000000000000000..e159a2cb2b5387fb24a9a2db8973f0994e3e319c --- /dev/null +++ b/scripts/enwik8/att_matching.sh @@ -0,0 +1,53 @@ +#NUMBER OF HEADS 4 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20/best \ + --data-path /mnt/data/vinhbk/datasets/enwik8 --dataset enwik8 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/enwik8 --dataset enwik8 + +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20/best \ + --model-b /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/enwik8 --dataset enwik8 + +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20/best \ + --data-path /mnt/data/vinhbk/datasets/enwik8 --dataset enwik8 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/enwik8 --dataset enwik8 + +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20/best \ + --model-b /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/enwik8 --dataset enwik8 + + + +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20/best \ + --data-path /mnt/data/vinhbk/datasets/enwik8 --dataset enwik8 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/enwik8 --dataset enwik8 + +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20/best \ + --model-b /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/enwik8 --dataset enwik8 diff --git a/scripts/enwik8/finetune.sh b/scripts/enwik8/finetune.sh new file mode 100644 index 0000000000000000000000000000000000000000..4619a9ab1c7dad5353baea3ce34a7bd71e9fd94e --- /dev/null +++ b/scripts/enwik8/finetune.sh @@ -0,0 +1,140 @@ +#LEARNABLE +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 4 --lmc-layer-indices 0 \ +# --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 4 --lmc-layer-indices 0 \ +# --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 4 --lmc-layer-indices 0 \ +# --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 8 --lmc-layer-indices 0 \ +# --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 8 --lmc-layer-indices 0 \ +# --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 \ +# --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 \ +# --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 \ +# --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 16 --lmc-layer-indices 0 \ +# --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +#ROPE +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-learnable-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +# --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 \ +# --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ +# --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-rope-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +# --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ +# --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ +# --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + diff --git a/scripts/enwik8/finetune4.sh b/scripts/enwik8/finetune4.sh new file mode 100644 index 0000000000000000000000000000000000000000..7ffd71b7914badfe81dfe7d487ef7cda9b2293af --- /dev/null +++ b/scripts/enwik8/finetune4.sh @@ -0,0 +1,57 @@ +WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 8 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 8 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 16 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 4 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 4 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 4 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 \ No newline at end of file diff --git a/scripts/enwik8/finetune5.sh b/scripts/enwik8/finetune5.sh new file mode 100644 index 0000000000000000000000000000000000000000..14f14110a918ceaf726986bb91b5297ed37d4a9c --- /dev/null +++ b/scripts/enwik8/finetune5.sh @@ -0,0 +1,57 @@ +WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ +# --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/enwik8/lr0.00025-rope-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_56000 \ + --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --finetune-mlp --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 /mnt/data/vinhbk/weights/enwik8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/enwik8 \ No newline at end of file diff --git a/scripts/enwik8/plot.sh b/scripts/enwik8/plot.sh new file mode 100644 index 0000000000000000000000000000000000000000..903cf19294c172491d8956c9b752c017c426549f --- /dev/null +++ b/scripts/enwik8/plot.sh @@ -0,0 +1,4 @@ +python src/lgmodeling/plot.py --output-dir plots/enwik8/learnable-indice0-heads4-shared1-routed0-topk0.pdf\ + --file-1 results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json\ + --file-2 results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json\ + --file-3 results/enwik8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json \ No newline at end of file diff --git a/scripts/enwik8/rope_matching.sh b/scripts/enwik8/rope_matching.sh new file mode 100644 index 0000000000000000000000000000000000000000..6f869a988b1b85da249581b4ce290b1fdc6b431b --- /dev/null +++ b/scripts/enwik8/rope_matching.sh @@ -0,0 +1,5 @@ +CUDA_VISIBLE_DEVICES=4 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-rope-indice0-heads8-shared1-routed0-topk0-seed0/best_60000 \ + --model-b /mnt/data/vinhbk/weights/enwik8/gpt2-finetune/finetune-rope-indice0-heads8-shared1-routed0-topk0-seed20/best_48000 \ + --data-path /mnt/data/vinhbk/datasets/enwik8 --dataset enwik8 \ No newline at end of file diff --git a/scripts/enwik8/train_model.sh b/scripts/enwik8/train_model.sh new file mode 100644 index 0000000000000000000000000000000000000000..0544e6dc7fadbe9e58f68d047a65b9662da371d6 --- /dev/null +++ b/scripts/enwik8/train_model.sh @@ -0,0 +1,66 @@ +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 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 enwik8 \ +# --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/enwik8 --data-path /mnt/data/vinhbk/datasets/enwik8 +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ +# --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "rope" \ +# --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 enwik8 \ +# --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/enwik8 --data-path /mnt/data/vinhbk/datasets/enwik8 +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ +# --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "sinusoidal" \ +# --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 enwik8 \ +# --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/enwik8 --data-path /mnt/data/vinhbk/datasets/enwik8 + + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "learnable" \ + --num-shared-experts 0 --num-routed-experts 4 --topk 4 --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 enwik8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/enwik8 --data-path /mnt/data/vinhbk/datasets/enwik8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "rope" \ + --num-shared-experts 0 --num-routed-experts 4 --topk 4 --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 enwik8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/enwik8 --data-path /mnt/data/vinhbk/datasets/enwik8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "sinusoidal" \ + --num-shared-experts 0 --num-routed-experts 4 --topk 4 --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 enwik8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/enwik8 --data-path /mnt/data/vinhbk/datasets/enwik8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "learnable" \ + --num-shared-experts 0 --num-routed-experts 4 --topk 2 --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 enwik8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-SMoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/enwik8 --data-path /mnt/data/vinhbk/datasets/enwik8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "rope" \ + --num-shared-experts 0 --num-routed-experts 4 --topk 2 --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 enwik8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-SMoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/enwik8 --data-path /mnt/data/vinhbk/datasets/enwik8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "sinusoidal" \ + --num-shared-experts 0 --num-routed-experts 4 --topk 2 --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 enwik8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Enwik8-SMoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/enwik8 --data-path /mnt/data/vinhbk/datasets/enwik8 \ No newline at end of file diff --git a/scripts/git.sh b/scripts/git.sh new file mode 100644 index 0000000000000000000000000000000000000000..4f02c05e6908bc04da21072ab28872e12a0e15cb --- /dev/null +++ b/scripts/git.sh @@ -0,0 +1,8 @@ +git add README.md +git add plots +git add results +git add scripts +git add src +git add pyproject.toml +git add requirements.txt +git add README.md diff --git a/scripts/imagenet/att_matching.sh b/scripts/imagenet/att_matching.sh new file mode 100644 index 0000000000000000000000000000000000000000..973de38aff7c3fdca8319f4a2f68fca1c257ab71 --- /dev/null +++ b/scripts/imagenet/att_matching.sh @@ -0,0 +1,40 @@ +# CUDA_VISIBLE_DEVICES=7 python src/imagenet/att_matching.py \ +# --model-a /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0/best_65052 \ +# --model-b /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20/best_5004 \ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-path plots/imagenet +# CUDA_VISIBLE_DEVICES=7 python src/imagenet/att_matching.py \ +# --model-a /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0/best_65052 \ +# --model-b /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40/best_5004 \ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-path plots/imagenet +# CUDA_VISIBLE_DEVICES=7 python src/imagenet/att_matching.py \ +# --model-a /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20/best_5004 \ +# --model-b /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40/best_5004 \ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-path plots/imagenet + + +CUDA_VISIBLE_DEVICES=7 python src/imagenet/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0/best_10008 \ + --model-b /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20/best_85068 \ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-path plots/imagenet +CUDA_VISIBLE_DEVICES=7 python src/imagenet/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0/best_10008 \ + --model-b /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40/best_35028 \ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-path plots/imagenet +CUDA_VISIBLE_DEVICES=7 python src/imagenet/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20/best_85068 \ + --model-b /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40/best_35028 \ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-path plots/imagenet + + +CUDA_VISIBLE_DEVICES=7 python src/imagenet/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0/best_5004 \ + --model-b /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20/best_5004 \ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-path plots/imagenet +CUDA_VISIBLE_DEVICES=7 python src/imagenet/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0/best_5004 \ + --model-b /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40/best_5004 \ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-path plots/imagenet +CUDA_VISIBLE_DEVICES=7 python src/imagenet/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20/best_5004 \ + --model-b /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40/best_5004 \ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-path plots/imagenet diff --git a/scripts/imagenet/finetune0.sh b/scripts/imagenet/finetune0.sh new file mode 100644 index 0000000000000000000000000000000000000000..6be91fac5864da0c1f3955638cc6c9506a0aeabb --- /dev/null +++ b/scripts/imagenet/finetune0.sh @@ -0,0 +1,33 @@ +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ +# --lr 0.0005 --epochs 100 --batch-size 256 --seed 0 --num-attention-heads 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ +# --lr 0.0005 --epochs 100 --batch-size 256 --seed 20 --num-attention-heads 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ +# --lr 0.0005 --epochs 100 --batch-size 256 --seed 40 --num-attention-heads 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ +# --lr 0.0005 --epochs 100 --batch-size 256 --seed 0 --num-attention-heads 16 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ + --lr 0.0005 --epochs 100 --batch-size 256 --seed 20 --num-attention-heads 16 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + + + diff --git a/scripts/imagenet/finetune1.sh b/scripts/imagenet/finetune1.sh new file mode 100644 index 0000000000000000000000000000000000000000..4743b4f5c3dc8a61aaebdd385ef805f5f731c910 --- /dev/null +++ b/scripts/imagenet/finetune1.sh @@ -0,0 +1,42 @@ +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ +# --lr 0.0005 --epochs 100 --batch-size 256 --seed 0 --num-attention-heads 12 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ +# --lr 0.0005 --epochs 100 --batch-size 256 --seed 20 --num-attention-heads 12 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ +# --lr 0.0005 --epochs 100 --batch-size 256 --seed 40 --num-attention-heads 12 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012\ +# --lr 0.0005 --epochs 100 --batch-size 256 --seed 0 --num-attention-heads 16 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 100 --batch-size 256 --seed 20 --num-attention-heads 16 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ + --lr 0.0005 --epochs 100 --batch-size 256 --seed 40 --num-attention-heads 16 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 100 --batch-size 256 --seed 40 --num-attention-heads 16 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + + diff --git a/scripts/imagenet/finetune2.sh b/scripts/imagenet/finetune2.sh new file mode 100644 index 0000000000000000000000000000000000000000..3383f0b154fe297b7540a147f88b8166bc947664 --- /dev/null +++ b/scripts/imagenet/finetune2.sh @@ -0,0 +1,58 @@ +<<<<<<< HEAD +#LEARNABLE +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2 python src/imagenet/test.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 12 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 20 --num-attention-heads 12 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 40 --num-attention-heads 12 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 8 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 20 --num-attention-heads 8 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ +# --lr 0.0005 --epochs 30 --batch-size 256 --seed 40 --num-attention-heads 8 --lmc-layer-indices 0\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ +# --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 16 --lmc-layer-indices 0\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ +# --lr 0.0005 --epochs 30 --batch-size 256 --seed 20 --num-attention-heads 16 --lmc-layer-indices 0\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ +# --lr 0.0005 --epochs 30 --batch-size 256 --seed 40 --num-attention-heads 16 --lmc-layer-indices 0\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + + +======= +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 100 --batch-size 256 --seed 20 --num-attention-heads 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +>>>>>>> master diff --git a/scripts/imagenet/finetune3.sh b/scripts/imagenet/finetune3.sh new file mode 100644 index 0000000000000000000000000000000000000000..7241a884eaac607487be2f00c025fead207efcd5 --- /dev/null +++ b/scripts/imagenet/finetune3.sh @@ -0,0 +1,54 @@ +<<<<<<< HEAD +#ROPE + +WANDB_MODE=offline CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_920736 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 12 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_920736\ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 20 --num-attention-heads 12 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_920736 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 40 --num-attention-heads 12 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_920736 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 8 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_920736 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 20 --num-attention-heads 8 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_920736 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 40 --num-attention-heads 8 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_920736 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 16 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_920736 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 20 --num-attention-heads 16 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + + +======= +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 100 --batch-size 256 --seed 40 --num-attention-heads 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +>>>>>>> master diff --git a/scripts/imagenet/finetune4.sh b/scripts/imagenet/finetune4.sh new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/imagenet/finetune6.sh b/scripts/imagenet/finetune6.sh new file mode 100644 index 0000000000000000000000000000000000000000..60228c50e8e7338d134a34c1b6afdb5a7012bfc9 --- /dev/null +++ b/scripts/imagenet/finetune6.sh @@ -0,0 +1,35 @@ +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ +# --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 12 --lmc-layer-indices 0\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 20 --num-attention-heads 12 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 40 --num-attention-heads 12 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 8 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 20 --num-attention-heads 8 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 40 --num-attention-heads 8 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune diff --git a/scripts/imagenet/finetune7.sh b/scripts/imagenet/finetune7.sh new file mode 100644 index 0000000000000000000000000000000000000000..d9d7ccc20fec44d7f1ab0b3b1ef255fde98d44e4 --- /dev/null +++ b/scripts/imagenet/finetune7.sh @@ -0,0 +1,35 @@ +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=7 python src/imagenet/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ +# --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 6 --lmc-layer-indices 0\ +# --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=7 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 20 --num-attention-heads 6 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=7 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 40 --num-attention-heads 6 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=7 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 16 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=7 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 20 --num-attention-heads 16 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=7 python src/imagenet/finetune.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/best_1266012 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 40 --num-attention-heads 16 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune \ No newline at end of file diff --git a/scripts/imagenet/plot.sh b/scripts/imagenet/plot.sh new file mode 100644 index 0000000000000000000000000000000000000000..97790081aa860ddbea02c30162ff4bef001e1737 --- /dev/null +++ b/scripts/imagenet/plot.sh @@ -0,0 +1,4 @@ +python src/imagenet/plot.py --output-dir plots/imagenet/\ + --file-1 results/imagenet/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20.json\ + --file-2 results/imagenet/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20.json\ + --file-3 results/imagenet/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20.json \ No newline at end of file diff --git a/scripts/imagenet/temp.sh b/scripts/imagenet/temp.sh new file mode 100644 index 0000000000000000000000000000000000000000..7c57fd074a7a89f8ed1a933e3301113272182095 --- /dev/null +++ b/scripts/imagenet/temp.sh @@ -0,0 +1,16 @@ +CUDA_VISIBLE_DEVICES=7 python src/imagenet/test.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-learnable-epochs300-batch256-shared1-routed0-topk0/best_870696 \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 12 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune +<<<<<<< HEAD +======= + + + +>>>>>>> master +CUDA_VISIBLE_DEVICES=6 python src/imagenet/test.py \ + --model-path /mnt/data/vinhbk/weights/imagenet/vit-finetune/finetune-learnable-indice[0]-heads12-shared1-routed0-topk0-seed0/best \ + --lr 0.0005 --epochs 30 --batch-size 256 --seed 0 --num-attention-heads 12 --lmc-layer-indices 0\ + --wandb-project LMC-Attention --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --data-path /mnt/data/vinhbk/datasets/imagenet --save-dir /mnt/data/vinhbk/weights/imagenet/vit-finetune \ No newline at end of file diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh new file mode 100644 index 0000000000000000000000000000000000000000..9b60e5931a33094fd279b97c61405930e03788d0 --- /dev/null +++ b/scripts/imagenet/train_model.sh @@ -0,0 +1,30 @@ +WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ +# --lr 0.0005 --epochs 300 --batch-size 256 \ +# --seed 0 --position-embeddings "sinusoidal"\ +# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ +# --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ +# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ +# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +# --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 diff --git a/scripts/lm1b/att_matching.sh b/scripts/lm1b/att_matching.sh new file mode 100644 index 0000000000000000000000000000000000000000..7da77496f7a8916d3e658221bf96c25060752a01 --- /dev/null +++ b/scripts/lm1b/att_matching.sh @@ -0,0 +1,46 @@ +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20/best \ + --data-path /mnt/data/vinhbk/datasets/lm1b --dataset lm1b +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/lm1b --dataset lm1b +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed20/best \ + --model-b /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/lm1b --dataset lm1b + +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20/best \ + --data-path /mnt/data/vinhbk/datasets/lm1b --dataset lm1b +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/lm1b --dataset lm1b +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20/best \ + --model-b /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/lm1b --dataset lm1b + +CUDA_VISIBLE_DEVICES=4 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0/best_56000 \ + --model-b /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20/best_96000 \ + --data-path /mnt/data/vinhbk/datasets/lm1b --dataset lm1b +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0/best_56000 \ + --model-b /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40/best_96000 \ + --data-path /mnt/data/vinhbk/datasets/lm1b --dataset lm1b +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20/best_96000 \ + --model-b /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40/best_96000 \ + --data-path /mnt/data/vinhbk/datasets/lm1b --dataset lm1b + + +# CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ +# --model-a /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0,1,2,3,4,5,6,7,8,9,10,11-heads12-shared1-routed0-topk0-mlpFalse-seed0/best_484000 \ +# --model-b /mnt/data/vinhbk/weights/lm1b/gpt2-finetune/finetune-learnable-indice0,1,2,3,4,5,6,7,8,9,10,11-heads12-shared1-routed0-topk0-mlpFalse-seed20/best_468000\ +# --data-path /mnt/data/vinhbk/datasets/lm1b --dataset lm1b + +# #best_476000 \ No newline at end of file diff --git a/scripts/lm1b/finetune_att.sh b/scripts/lm1b/finetune_att.sh new file mode 100644 index 0000000000000000000000000000000000000000..c418ff547386310b62f7bcb7b07c78a092ebc5b2 --- /dev/null +++ b/scripts/lm1b/finetune_att.sh @@ -0,0 +1,42 @@ +#LEARNABLE +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 4 --lmc-layer-indices 0 \ +# --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ +# --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ +# --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ +# --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ +# --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b +# ROPE +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-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 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b \ No newline at end of file diff --git a/scripts/lm1b/finetune_att0.sh b/scripts/lm1b/finetune_att0.sh new file mode 100644 index 0000000000000000000000000000000000000000..e5a98e4b9e295b018b431eea4f407ddcfc11f877 --- /dev/null +++ b/scripts/lm1b/finetune_att0.sh @@ -0,0 +1,84 @@ +#LEARNABLE +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 12 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 12 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b +#ROPE +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-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 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 12 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 12 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-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 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b \ No newline at end of file diff --git a/scripts/lm1b/finetune_att01.sh b/scripts/lm1b/finetune_att01.sh new file mode 100644 index 0000000000000000000000000000000000000000..d0af3f3717792694630d49480ed435a1f45662ae --- /dev/null +++ b/scripts/lm1b/finetune_att01.sh @@ -0,0 +1,57 @@ +# NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-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 16 --lmc-layer-indices 0 \ +# --learning-rate 0.00025 --batch-size 96 --max_step 200000 --warmup_step 1000 --dataset lm1b \ +# --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +# NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ +# --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 16 --lmc-layer-indices 0 \ +# --learning-rate 0.00025 --batch-size 96 --max_step 200000 --warmup_step 1000 --dataset lm1b \ +# --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +# NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ +# --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 16 --lmc-layer-indices 0 \ +# --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +# NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-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 500000 --warmup_step 2000 --dataset lm1b --finetune-mlp\ +# --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + + +# NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ +# --seed 20 --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 500000 --warmup_step 2000 --dataset lm1b --finetune-mlp\ +# --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + + +# NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ +# --seed 40 --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 500000 --warmup_step 2000 --dataset lm1b --finetune-mlp\ +# --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 500000 --warmup_step 2000 --dataset lm1b --finetune-mlp\ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --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 500000 --warmup_step 2000 --dataset lm1b --finetune-mlp\ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b diff --git a/scripts/lm1b/finetune_att1.sh b/scripts/lm1b/finetune_att1.sh new file mode 100644 index 0000000000000000000000000000000000000000..289ab760fa885442dd3864eea029757720fd59c4 --- /dev/null +++ b/scripts/lm1b/finetune_att1.sh @@ -0,0 +1,63 @@ +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1,2 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1,2 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --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 200000 --warmup_step 1000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b Xin chào mọi người tui te + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1,2 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --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 200000 --warmup_step 1000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1,2 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ + --learning-rate 0.00025 --batch-size 96 --max_step 200000 --warmup_step 1000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1,2 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ + --learning-rate 0.00025 --batch-size 96 --max_step 200000 --warmup_step 1000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1,2 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ + --learning-rate 0.00025 --batch-size 96 --max_step 200000 --warmup_step 1000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-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 8 --lmc-layer-indices 0 \ + --learning-rate 0.00025 --batch-size 96 --max_step 200000 --warmup_step 1000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 8 --lmc-layer-indices 0 \ + --learning-rate 0.00025 --batch-size 96 --max_step 200000 --warmup_step 1000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 8 --lmc-layer-indices 0 \ + --learning-rate 0.00025 --batch-size 96 --max_step 200000 --warmup_step 1000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b \ No newline at end of file diff --git a/scripts/lm1b/finetune_att2.sh b/scripts/lm1b/finetune_att2.sh new file mode 100644 index 0000000000000000000000000000000000000000..71c454bc5c2f2fc99e3c3a74aff5db4fb17a9a27 --- /dev/null +++ b/scripts/lm1b/finetune_att2.sh @@ -0,0 +1,42 @@ +#LEARNABLE +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3,5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 16 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3,5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 16 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3,5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 16 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b +# ROPE +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3,5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-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 16 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3,5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 16 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=3,5 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 16 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b \ No newline at end of file diff --git a/scripts/lm1b/finetune_att23.sh b/scripts/lm1b/finetune_att23.sh new file mode 100644 index 0000000000000000000000000000000000000000..181a93080f33f60ae1c0cf2b85975c5c9fb98b63 --- /dev/null +++ b/scripts/lm1b/finetune_att23.sh @@ -0,0 +1,52 @@ +# # #LEARNABLE +# NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ +# --seed 20 --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 500000 --warmup_step 2000 --dataset lm1b \ +# --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +# NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/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 500000 --warmup_step 2000 --dataset lm1b \ +# --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +# NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ +# --seed 40 --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 500000 --warmup_step 2000 --dataset lm1b \ +# --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-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 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --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 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --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 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + + +NVIDIA_TF32_OVERRIDE=0 JAX_DEFAULT_MATMUL_PRECISION=highest CUBLAS_WORKSPACE_CONFIG=:4096:8 WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --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 500000 --warmup_step 2000 --dataset lm1b --finetune-mlp\ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + + diff --git a/scripts/lm1b/finetune_att4.sh b/scripts/lm1b/finetune_att4.sh new file mode 100644 index 0000000000000000000000000000000000000000..875046db393b278df26ee06e121e86e71d4f59b1 --- /dev/null +++ b/scripts/lm1b/finetune_att4.sh @@ -0,0 +1,42 @@ +#LEARNABLE +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b +# ROPE +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-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 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,4 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b \ No newline at end of file diff --git a/scripts/lm1b/finetune_att67.sh b/scripts/lm1b/finetune_att67.sh new file mode 100644 index 0000000000000000000000000000000000000000..e67108d09a7a7055e95ccb01cdc07e70c0f5cb9a --- /dev/null +++ b/scripts/lm1b/finetune_att67.sh @@ -0,0 +1,55 @@ +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=6,7 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ +# --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 16 --lmc-layer-indices 0 \ +# --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6,7 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6,7 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/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 6 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6,7 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 6 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6,7 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-learnable-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 6 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6,7 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-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 6 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6,7 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 6 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6,7 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/lm1b/lr0.00025-rope-step500000-warm2000-size96-layer12-embd768-heads12-shared1-routed0-topk0/best_476000 \ + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 6 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/lm1b/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/lm1b \ No newline at end of file diff --git a/scripts/lm1b/plot.sh b/scripts/lm1b/plot.sh new file mode 100644 index 0000000000000000000000000000000000000000..32c58c1b17ecadf5661a8747929b3c242e4bdb22 --- /dev/null +++ b/scripts/lm1b/plot.sh @@ -0,0 +1,4 @@ +python src/lgmodeling/plot.py --output-dir plots/lm1b/learnable-indice0-heads16-shared1-routed0-topk0.pdf\ + --file-1 results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20].json\ + --file-2 results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json\ + --file-3 results/lm1b/[finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40].json \ No newline at end of file diff --git a/scripts/lm1b/train_model.sh b/scripts/lm1b/train_model.sh new file mode 100644 index 0000000000000000000000000000000000000000..3f16ac8d9d7d11e97d32772b4cb6284cd64915be --- /dev/null +++ b/scripts/lm1b/train_model.sh @@ -0,0 +1,77 @@ +<<<<<<< HEAD +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/train_model.py \ +======= +WANDB_MODE=offline CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/train_model.py \ +>>>>>>> master + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "rope" \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0 --rotary-dim 64 \ + --n_layer 12 --n_embd 768 --n_head 12 --n_inner 3072 --attention-bias \ + --learning-rate 0.00025 --batch-size 96 --max_step 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "sinusoidal" \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0 --rotary-dim 64 \ + --n_layer 12 --n_embd 768 --n_head 12 --n_inner 3072 --attention-bias \ + --learning-rate 0.00025 --batch-size 96 --max_step 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "" \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0 --rotary-dim 64 \ + --n_layer 12 --n_embd 768 --n_head 12 --n_inner 3072 --attention-bias \ + --learning-rate 0.00025 --batch-size 96 --max_step 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b --data-path /mnt/data/vinhbk/datasets/lm1b + + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "rope" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 4 --rotary-dim 64 \ + --n_layer 12 --n_embd 768 --n_head 12 --n_inner 3072 --attention-bias \ + --learning-rate 0.00025 --batch-size 96 --max_step 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "sinusoidal" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 4 --rotary-dim 64 \ + --n_layer 12 --n_embd 768 --n_head 12 --n_inner 3072 --attention-bias \ + --learning-rate 0.00025 --batch-size 96 --max_step 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 4 --rotary-dim 64 \ + --n_layer 12 --n_embd 768 --n_head 12 --n_inner 3072 --attention-bias \ + --learning-rate 0.00025 --batch-size 96 --max_step 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b --data-path /mnt/data/vinhbk/datasets/lm1b + + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=0,1 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "rope" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 2 --rotary-dim 64 \ + --n_layer 12 --n_embd 768 --n_head 12 --n_inner 3072 --attention-bias \ + --learning-rate 0.00025 --batch-size 96 --max_step 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "sinusoidal" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 2 --rotary-dim 64 \ + --n_layer 12 --n_embd 768 --n_head 12 --n_inner 3072 --attention-bias \ + --learning-rate 0.00025 --batch-size 96 --max_step 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b --data-path /mnt/data/vinhbk/datasets/lm1b + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 2 --rotary-dim 64 \ + --n_layer 12 --n_embd 768 --n_head 12 --n_inner 3072 --attention-bias \ + --learning-rate 0.00025 --batch-size 96 --max_step 500000 --warmup_step 2000 --dataset lm1b \ + --wandb-project LMC-Attention --wandb-group "GPT2-OneBillionWord-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/lm1b --data-path /mnt/data/vinhbk/datasets/lm1b \ No newline at end of file diff --git a/scripts/modal.sh b/scripts/modal.sh new file mode 100644 index 0000000000000000000000000000000000000000..f0278cfdd8f7e36653c8e67ad3abda4dde5c8519 --- /dev/null +++ b/scripts/modal.sh @@ -0,0 +1,6 @@ +modal volume create weights +modal volume create datasets +modal volume create plots +modal volume create results + + diff --git a/scripts/text8/att_matching.sh b/scripts/text8/att_matching.sh new file mode 100644 index 0000000000000000000000000000000000000000..84cb8e3280c1fe46920322fceeff21f6d3f77726 --- /dev/null +++ b/scripts/text8/att_matching.sh @@ -0,0 +1,50 @@ +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0/best_60000\ + --model-b /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20/best_60000 \ + --data-path /mnt/data/vinhbk/datasets/text8 --dataset text8 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0/best_60000\ + --model-b /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40/best_60000 \ + --data-path /mnt/data/vinhbk/datasets/text8 --dataset text8 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20/best_60000\ + --model-b /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40/best_60000 \ + --data-path /mnt/data/vinhbk/datasets/text8 --dataset text8 + + + +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0/best_60000\ + --model-b /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20/best_60000 \ + --data-path /mnt/data/vinhbk/datasets/text8 --dataset text8 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed0/best_60000\ + --model-b /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40/best_60000 \ + --data-path /mnt/data/vinhbk/datasets/text8 --dataset text8 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed20/best_60000\ + --model-b /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads8-shared1-routed0-topk0-seed40/best_60000 \ + --data-path /mnt/data/vinhbk/datasets/text8 --dataset text8 + + +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0/best_60000\ + --model-b /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20/best_60000 \ + --data-path /mnt/data/vinhbk/datasets/text8 --dataset text8 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed0/best_60000\ + --model-b /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40/best_60000 \ + --data-path /mnt/data/vinhbk/datasets/text8 --dataset text8 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed20/best_60000\ + --model-b /mnt/data/vinhbk/weights/text8/gpt2-finetune/finetune-learnable-indice0-heads16-shared1-routed0-topk0-seed40/best_60000 \ + --data-path /mnt/data/vinhbk/datasets/text8 --dataset text8 \ No newline at end of file diff --git a/scripts/text8/finetune.sh b/scripts/text8/finetune.sh new file mode 100644 index 0000000000000000000000000000000000000000..131ea7b8aa53c796cb36e35c28a92724ccc422d8 --- /dev/null +++ b/scripts/text8/finetune.sh @@ -0,0 +1,117 @@ +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ + --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ + --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --lmc-layer-indices 0\ + --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ + --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ + --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --lmc-layer-indices 0\ + --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ + --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ + --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --lmc-layer-indices 0 \ + --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 + + + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ +# --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ +# --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ +# --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ +# --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ +# --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ +# --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ +# --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ +# --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=1 python src/lgmodeling/finetune.py \ +# --model-path /mnt/data/vinhbk/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0 \ +# --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 /mnt/data/vinhbk/weights/text8/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/text8 diff --git a/scripts/text8/plot.sh b/scripts/text8/plot.sh new file mode 100644 index 0000000000000000000000000000000000000000..2583b644861d5e3a20a2296d0e7200d041588654 --- /dev/null +++ b/scripts/text8/plot.sh @@ -0,0 +1,4 @@ +python src/lgmodeling/plot.py --output-dir plots/text8/learnable-indice0-heads4-shared1-routed0-topk0.pdf\ + --file-1 results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json\ + --file-2 results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json\ + --file-3 results/text8/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json \ No newline at end of file diff --git a/scripts/text8/train_model.sh b/scripts/text8/train_model.sh new file mode 100644 index 0000000000000000000000000000000000000000..fefb211976ff095463569c14339efe30b283cf57 --- /dev/null +++ b/scripts/text8/train_model.sh @@ -0,0 +1,70 @@ +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 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 72000 --warmup_step 2000 --dataset text8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /cm/archive/vinhbk1/weights/text8 --data-path /cm/archive/vinhbk1/datasets/text8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "rope" \ + --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 72000 --warmup_step 2000 --dataset text8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /cm/archive/vinhbk1/weights/text8 --data-path /cm/archive/vinhbk1/datasets/text8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "sinusoidal" \ + --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 72000 --warmup_step 2000 --dataset text8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /cm/archive/vinhbk1/weights/text8 --data-path /cm/archive/vinhbk1/datasets/text8 + +<<<<<<< HEAD + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 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 4 --topk 4 --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 72000 --warmup_step 2000 --dataset text8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Text8-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /cm/archive/vinhbk1/weights/text8 --data-path /cm/archive/vinhbk1/datasets/text8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "rope" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 4 --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 72000 --warmup_step 2000 --dataset text8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Text8-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /cm/archive/vinhbk1/weights/text8 --data-path /cm/archive/vinhbk1/datasets/text8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "sinusoidal" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 4 --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 72000 --warmup_step 2000 --dataset text8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Text8-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /cm/archive/vinhbk1/weights/text8 --data-path /cm/archive/vinhbk1/datasets/text8 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 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 4 --topk 2 --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 72000 --warmup_step 2000 --dataset text8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Text8-SMoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /cm/archive/vinhbk1/weights/text8 --data-path /cm/archive/vinhbk1/datasets/text8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "rope" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 2 --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 72000 --warmup_step 2000 --dataset text8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Text8-SMoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /cm/archive/vinhbk1/weights/text8 --data-path /cm/archive/vinhbk1/datasets/text8 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --position-embeddings "sinusoidal" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 2 --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 72000 --warmup_step 2000 --dataset text8 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Text8-SMoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /cm/archive/vinhbk1/weights/text8 --data-path /cm/archive/vinhbk1/datasets/text8 +======= + +>>>>>>> master diff --git a/scripts/wt103/att_matching.sh b/scripts/wt103/att_matching.sh new file mode 100644 index 0000000000000000000000000000000000000000..b78a3511b5a61d72cb90842aeb06f54e8a695909 --- /dev/null +++ b/scripts/wt103/att_matching.sh @@ -0,0 +1,38 @@ +# CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ +# --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0/best \ +# --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed20/best \ +# --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 +# CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ +# --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0/best \ +# --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed40/best \ +# --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 +# CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ +# --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed20/best \ +# --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed40/best \ +# --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 + +# CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ +# --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed0/best \ +# --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed20/best \ +# --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed20/best \ + --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads3-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 + +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20/best \ + --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20/best \ + --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 diff --git a/scripts/wt103/finetune_att6.sh b/scripts/wt103/finetune_att6.sh new file mode 100644 index 0000000000000000000000000000000000000000..1cd25d8a1fbcdef55719c035c7280e959aad1140 --- /dev/null +++ b/scripts/wt103/finetune_att6.sh @@ -0,0 +1,180 @@ +#LEARNABLE +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-learnable-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 \ +======= + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-learnable-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 \ +======= + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-learnable-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 \ +======= + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-learnable-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 \ +======= + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-learnable-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 \ +======= + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-learnable-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ +======= + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-learnable-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ +======= + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-learnable-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ +======= + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + +#ROPE +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-rope-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 \ +======= + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-rope-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 \ +======= + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-rope-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 \ +======= + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-rope-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 \ +======= + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-rope-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 \ +======= + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-rope-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 \ +======= + --seed 40 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 2 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-rope-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ +======= + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + + +WANDB_MODE=online CUDA_VISIBLE_DEVICES=6 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-rope-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ +<<<<<<< HEAD + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 \ +======= + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 4 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11 \ +>>>>>>> master + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 + diff --git a/scripts/wt103/finetune_att7.sh b/scripts/wt103/finetune_att7.sh new file mode 100644 index 0000000000000000000000000000000000000000..3cd1f5d7a88efa07b40df3e2dbc1cdad9895ae36 --- /dev/null +++ b/scripts/wt103/finetune_att7.sh @@ -0,0 +1,6 @@ +WANDB_MODE=online CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/finetune.py \ + --model-path /mnt/data/vinhbk/weights/wt103/lr0.00025-learnable-step60000-warm2000-size64-layer12-embd192-heads3-shared1-routed0-topk0/best_60000 \ + --seed 20 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --n_head 3 --lmc-layer-indices 0 \ + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103/gpt2-finetune --data-path /mnt/data/vinhbk/datasets/wt103 diff --git a/scripts/wt103/plot.sh b/scripts/wt103/plot.sh new file mode 100644 index 0000000000000000000000000000000000000000..dcdd218bbf58e4d3b52a3ab6ba6fd304be0f2840 --- /dev/null +++ b/scripts/wt103/plot.sh @@ -0,0 +1,4 @@ +python src/lgmodeling/plot.py --output-dir plots/wt103/learnable-indice0-heads4-shared1-routed0-topk0.pdf\ + --file-1 results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20].json\ + --file-2 results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed0+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json\ + --file-3 results/wt103/[finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed20+finetune-learnable-indice0-heads4-shared1-routed0-topk0-seed40].json \ No newline at end of file diff --git a/scripts/wt103/rope_matching.sh b/scripts/wt103/rope_matching.sh new file mode 100644 index 0000000000000000000000000000000000000000..4a4c7b8bb57342c886ba16b8316fdedca67b45d9 --- /dev/null +++ b/scripts/wt103/rope_matching.sh @@ -0,0 +1,30 @@ +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-rope-indice0-heads3-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-rope-indice0-heads3-shared1-routed0-topk0-seed20/best \ + --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 + +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed20/best \ + --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 + +CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-learnable-indice0-heads2-shared1-routed0-topk0-seed40/best \ + --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 + +CUDA_VISIBLE_DEVICES=4 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-rope-indice0-heads3-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-rope-indice0-heads3-shared1-routed0-topk0-seed20/best \ + --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 + + +CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/att_matching.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 \ + --model-a /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-rope-indice0-heads2-shared1-routed0-topk0-seed0/best \ + --model-b /mnt/data/vinhbk/weights/wt103/gpt2-finetune/finetune-rope-indice0-heads2-shared1-routed0-topk0-seed20/best \ + --data-path /mnt/data/vinhbk/datasets/wt103 --dataset wt103 \ No newline at end of file diff --git a/scripts/wt103/train_model.sh b/scripts/wt103/train_model.sh new file mode 100644 index 0000000000000000000000000000000000000000..6faa1f9cc91e3dc0fd4ed521ae7434521b364888 --- /dev/null +++ b/scripts/wt103/train_model.sh @@ -0,0 +1,46 @@ + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/train_model.py \ +# --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "learnable" \ +# --num-shared-experts 1 --num-routed-experts 4 --topk 4 --rotary-dim 64 \ +# --n_layer 12 --n_embd 192 --n_head 3 --n_inner 768 --attention-bias \ +# --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ +# --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-MoE" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/wt103 --data-path /mnt/data/vinhbk/datasets/wt103 + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/train_model.py \ +# --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "rope" \ +# --num-shared-experts 1 --num-routed-experts 4 --topk 4 --rotary-dim 64 \ +# --n_layer 12 --n_embd 192 --n_head 3 --n_inner 768 --attention-bias \ +# --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ +# --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-MoE" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/wt103 --data-path /mnt/data/vinhbk/datasets/wt103 +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/train_model.py \ +# --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "sinusoidal" \ +# --num-shared-experts 1 --num-routed-experts 4 --topk 4 --rotary-dim 64 \ +# --n_layer 12 --n_embd 192 --n_head 3 --n_inner 768 --attention-bias \ +# --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ +# --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-MoE" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/wt103 --data-path /mnt/data/vinhbk/datasets/wt103 + + +# WANDB_MODE=online CUDA_VISIBLE_DEVICES=5 python src/lgmodeling/train_model.py \ +# --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "learnable" \ +# --num-shared-experts 1 --num-routed-experts 4 --topk 2 --rotary-dim 64 \ +# --n_layer 12 --n_embd 192 --n_head 3 --n_inner 768 --attention-bias \ +# --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ +# --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-SMoE" --wandb-entity "vinh-bui0512-hcmut"\ +# --model-save-dir /mnt/data/vinhbk/weights/wt103 --data-path /mnt/data/vinhbk/datasets/wt103 +WANDB_MODE=online CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "rope" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 2 --rotary-dim 64 \ + --n_layer 12 --n_embd 192 --n_head 3 --n_inner 768 --attention-bias \ + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-SMoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103 --data-path /mnt/data/vinhbk/datasets/wt103 +WANDB_MODE=offline CUDA_VISIBLE_DEVICES=7 python src/lgmodeling/train_model.py \ + --seed 0 --tgt_len 256 --mem_len 256 --eval_tgt_len 256 --position-embeddings "sinusoidal" \ + --num-shared-experts 1 --num-routed-experts 4 --topk 2 --rotary-dim 64 \ + --n_layer 12 --n_embd 192 --n_head 3 --n_inner 768 --attention-bias \ + --learning-rate 0.00025 --batch-size 64 --max_step 60000 --warmup_step 2000 --dataset wt103 \ + --wandb-project LMC-Attention --wandb-group "GPT2-Wikitext103-SMoE" --wandb-entity "vinh-bui0512-hcmut"\ + --model-save-dir /mnt/data/vinhbk/weights/wt103 --data-path /mnt/data/vinhbk/datasets/wt103 diff --git a/src/__pycache__/datasets.cpython-310.pyc b/src/__pycache__/datasets.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..855a7a6ad459df5ac1020416f9385c67c694e03f Binary files /dev/null and b/src/__pycache__/datasets.cpython-310.pyc differ diff --git a/src/__pycache__/utils.cpython-310.pyc b/src/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c1ed34d74592a10bca5b39209952b74d1b8b070b Binary files /dev/null and b/src/__pycache__/utils.cpython-310.pyc differ diff --git a/src/datasets.py b/src/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..3bb26c2e8ecd450cfc0dfbbb10c7cebe641141b8 --- /dev/null +++ b/src/datasets.py @@ -0,0 +1,57 @@ +import numpy as np +import tensorflow_datasets as tfds +from torchvision import datasets, transforms +import torch, os +import pickle + +def load_cifar10(data_dir=None): + """Return the training and test datasets, as jnp.array's.""" + train_ds_images_u8, train_ds_labels = tfds.as_numpy( + tfds.load("cifar10", split="train", batch_size=-1, as_supervised=True, data_dir=data_dir)) + test_ds_images_u8, test_ds_labels = tfds.as_numpy( + tfds.load("cifar10", split="test", batch_size=-1, as_supervised=True, data_dir=data_dir)) + train_ds = {"images_u8": train_ds_images_u8, "labels": train_ds_labels} + test_ds = {"images_u8": test_ds_images_u8, "labels": test_ds_labels} + return train_ds, test_ds + +import tensorflow_datasets as tfds + +def load_cifar100(data_dir=None): + # Define paths for saved dataset + train_pickle = os.path.join(data_dir, "cifar100_train.pkl") + test_pickle = os.path.join(data_dir, "cifar100_test.pkl") + + # Check if dataset already exists + if os.path.exists(train_pickle) and os.path.exists(test_pickle): + print(f"Loading dataset from {data_dir}") + with open(train_pickle, "rb") as f: + train_ds = pickle.load(f) + with open(test_pickle, "rb") as f: + test_ds = pickle.load(f) + return train_ds, test_ds + + os.makedirs(data_dir, exist_ok=True) + + # Load CIFAR-100 dataset with optional data_dir + train_ds_images_u8, train_ds_labels = tfds.as_numpy( + tfds.load("cifar100", split="train", batch_size=-1, as_supervised=True, data_dir=data_dir)) + test_ds_images_u8, test_ds_labels = tfds.as_numpy( + tfds.load("cifar100", split="test", batch_size=-1, as_supervised=True, data_dir=data_dir)) + + # Organize datasets into dictionaries + train_ds = {"images_u8": train_ds_images_u8, "labels": train_ds_labels} + test_ds = {"images_u8": test_ds_images_u8, "labels": test_ds_labels} + + # Save datasets to pickle files + print(f"Saving dataset to {data_dir}") + try: + with open(train_pickle, "wb") as f: + pickle.dump(train_ds, f) + with open(test_pickle, "wb") as f: + pickle.dump(test_ds, f) + print(f"Successfully saved datasets to {data_dir}") + except Exception as e: + print(f"Error saving datasets: {e}") + raise + + return train_ds, test_ds \ No newline at end of file diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ed5568e24c9cb2a9cdc242d62f43c856d34f902e Binary files /dev/null and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ diff --git a/src/imagenet/__pycache__/engine.cpython-310.pyc b/src/imagenet/__pycache__/engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..13f4bd09c9fc87a89d2cd2771cf03aa0135ad245 Binary files /dev/null and b/src/imagenet/__pycache__/engine.cpython-310.pyc differ diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e21f534325d6fea4b6fe1bace318c98626520991 Binary files /dev/null and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ diff --git a/src/imagenet/__pycache__/matching_utils.cpython-310.pyc b/src/imagenet/__pycache__/matching_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b2761d002d29c753530b2dff6873e931e658f7cb Binary files /dev/null and b/src/imagenet/__pycache__/matching_utils.cpython-310.pyc differ diff --git a/src/imagenet/__pycache__/model.cpython-310.pyc b/src/imagenet/__pycache__/model.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40066d7e684c88122281c181eb6e6e3efce6798b Binary files /dev/null and b/src/imagenet/__pycache__/model.cpython-310.pyc differ diff --git a/src/imagenet/__pycache__/utils.cpython-310.pyc b/src/imagenet/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb433ad23364c23edcb1af3cb90b12c9f567e47f Binary files /dev/null and b/src/imagenet/__pycache__/utils.cpython-310.pyc differ diff --git a/src/imagenet/att_matching.py b/src/imagenet/att_matching.py new file mode 100644 index 0000000000000000000000000000000000000000..e2c238b7675765dae694662f203d7025b2c51cf9 --- /dev/null +++ b/src/imagenet/att_matching.py @@ -0,0 +1,196 @@ +import argparse +import os +import jax +import flax +import copy +import torch +import wandb +import optax +import itertools +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from jax import random, vmap +from flax import linen as nn +from typing import Any, Dict, List +from jax import random, vmap +from flax.jax_utils import replicate, unreplicate +from flax.core.frozen_dict import freeze, unfreeze +from flax.training import checkpoints, train_state +from flax.traverse_util import flatten_dict, unflatten_dict +from flax.training.common_utils import get_metrics, onehot, shard +from utils import flatten_params, lerp, unflatten_params +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from transformers.models.vit.modeling_flax_vit import ViTConfig +from matching_utils import weight_matching_attn +from datasets import build_dataset +from flax.serialization import from_bytes +import multiprocessing as mp +import matplotlib.pyplot as plt +import numpy as np +import json +mp.set_start_method("spawn", force=True) +def load_flax_params(checkpoint_dir, target): + msgpack_path = os.path.join(checkpoint_dir, "flax_model.msgpack") + with open(msgpack_path, "rb") as f: + packed_bytes = f.read() + params = from_bytes(target, packed_bytes) + return params +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val + +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def make_stuff(model): + apply_fn = model.__call__ + def batch_eval(params, batch): + outputs = apply_fn(params=params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_batch_eval = jax.pmap(batch_eval, axis_name="batch") + def dataset_loss_and_acc(params, dataloader): + """Evaluate over the dataloader on all devices""" + eval_results = [] + pbar = tqdm(dataloader, desc="Evaluating", leave=False) + rep_params = replicate(params) + for step_idx, (images, labels) in enumerate(pbar): + batch = prepare_image_batch(images, labels) + eval_metric = parallel_batch_eval(rep_params, batch) + eval_results.append(copy.deepcopy(eval_metric)) + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + loss, acc1 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]) + return loss, acc1 + return {"batch_eval": batch_eval, "dataset_loss_and_acc": dataset_loss_and_acc} +def compute_interpolation(params_a, params_b_target, lambdas, stuff, test_loader, desc="Interpolation"): + train_loss_interp, test_loss_interp = [], [] + train_acc_interp, test_acc_interp = [], [] + for lam in tqdm(lambdas, desc=desc): + p_interp = freeze(lerp(lam, unfreeze(params_a), unfreeze(params_b_target))) + # if lam == 0: with open("output.txt", "a") as f: f.write(f"{p_interp}\n") + test_loss, test_acc = stuff["dataset_loss_and_acc"](p_interp, test_loader) + test_loss_interp.append(test_loss) + test_acc_interp.append(test_acc) + return {"Test Loss": [float(f"{x:.4f}") for x in test_loss_interp],"Test Acc": [float(f"{x:.4f}") for x in test_acc_interp]} +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-a", type=str, required=True, help="Path to first fine-tuned GPT2 model checkpoint") + parser.add_argument("--model-b", type=str, required=True, help="Path to second fine-tuned GPT@ model checkpoint") + parser.add_argument("--batch-size", type=int, default=256) + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR10','CIFAR100', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=8) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--save-path", type=str, default="/", help="Path to plot directory") + args = parser.parse_args() + config = ViTConfig.from_json_file(os.path.join(os.path.dirname(args.model_a).rstrip("/"),'config.json')) + lmc_config = ViTConfig(**config.lmc_config) + config.lmc_config = lmc_config + model = LMCFlaxViTForImageClassification(config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype),) + print_model(model.params) + if os.path.exists(args.model_a) and os.path.exists(args.model_b): + params_a = load_flax_params(args.model_a,copy.deepcopy(model.params)) + params_b = load_flax_params(args.model_b,copy.deepcopy(model.params)) + # params_a = checkpoints.restore_checkpoint(ckpt_dir=args.model_a, target={"params": model.params})["params"] + # params_b = checkpoints.restore_checkpoint(ckpt_dir=args.model_b, target={"params": model.params})["params"] + else: + raise FileNotFoundError(f"Checkpoint path does not exist") + stuff = make_stuff(model = model) + train_loader, val_loader = imagenet_data_loader(args) + lambdas = jnp.linspace(0, 1, num=3) + rng = random.PRNGKey(args.seed) + # Compute naive interpolation + activation = None #stuff["get_attention_inputs"](params_b, val_loader) # None # + naive_results = compute_interpolation(copy.deepcopy(params_a), copy.deepcopy(params_b), lambdas, stuff, val_loader, desc="Naive Interpolation") + print(json.dumps({"Naive": naive_results}, indent=2)) + naive_results = compute_interpolation(copy.deepcopy(params_a), copy.deepcopy(params_b), lambdas, stuff, val_loader, desc="Naive Interpolation") + print(json.dumps({"Naive": naive_results}, indent=2)) + all_results = {"Naive": naive_results} + # Compute weight matching interpolations for each method + aligned_models = weight_matching_attn(rng, copy.deepcopy(params_a), copy.deepcopy(params_b), activation, config) + for method, params_b_aligned in aligned_models.items(): + method_results = compute_interpolation(params_a, params_b_aligned, lambdas, stuff, val_loader, desc=f"{method} Interpolation") + all_results[method] = method_results + print(json.dumps({method: method_results}, indent=2)) + # Save directories + os.makedirs(f"./plots/imagenet", exist_ok=True) + os.makedirs(f"./results/imagenet", exist_ok=True) + # Save results JSON + print("Save List of Values...") + name_a = os.path.basename(os.path.dirname(args.model_a).rstrip("/")) + name_b = os.path.basename(os.path.dirname(args.model_b).rstrip("/")) + result_path = f'results/imagenet/{name_a}+{name_b}.json' + with open(result_path, 'w') as f: + json.dump(all_results, f, indent=2) + # Plot + print("Generating plots...") + plot_path = f"./plots/imagenet/{name_a}+{name_b}.pdf" + plt.rcParams.update({ + "font.family": "serif", + 'legend.frameon': False, + 'lines.linewidth': 2, + 'font.size': 13, + 'axes.labelsize': 16, + 'xtick.labelsize': 11, + 'ytick.labelsize': 11, + 'legend.fontsize': 11, + }) + plt.style.use('tableau-colorblind10') + num_points = len(all_results["Naive"]["Test Loss"]) + lambda_values = np.linspace(0, 1, num_points) + fig, axs = plt.subplots(2, 2, figsize=(12, 10)) + metrics = ["Test Loss", "Test Acc", "Test Loss", "Test Acc",] + positions = [(0, 0), (0, 1), (1, 0), (1, 1)] + for metric, pos in zip(metrics, positions): + row, col = pos + ax = axs[row, col] + for method in all_results: + ax.plot(lambda_values, all_results[method][metric], label=method) + ax.set_xticks([0, 0.5, 1]) + ax.set_xticklabels(["Model 1", r"$\lambda$", "Model 2"]) + ax.set_ylabel(metric) + ax.legend(loc='best') + plt.tight_layout() + plt.savefig(plot_path) + plt.close() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/imagenet/data.sh b/src/imagenet/data.sh new file mode 100644 index 0000000000000000000000000000000000000000..13eb1144759a2f2710276dbc1ea55cf0eb9bf4bd --- /dev/null +++ b/src/imagenet/data.sh @@ -0,0 +1,28 @@ +echo "- Downloading ImageNet 2012 dataset (Imagenet 1k)" + +if [[ ! -d 'imagenet' ]]; then + mkdir -p imagenet + cd imagenet + echo "- Downloading ILSVRC2012_img_train.tar" + wget --no-check-certificate https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_train.tar + echo "- Downloading ILSVRC2012_img_val.tar" + wget --no-check-certificate https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar + echo "- Extracting validation set" + mkdir -p val + tar -xvf ILSVRC2012_img_val.tar -C val + cd val + wget https://raw.githubusercontent.com/PatrickHua/EasyImageNet/main/valprep.sh + sh valprep.sh + cd .. + echo "- Extracting training set" + mkdir -p train + tar -xvf ILSVRC2012_img_train.tar -C train + cd train + find . -name "*.tar" | while read NAME ; do + mkdir -p "${NAME%.tar}" + tar -xvf "${NAME}" -C "${NAME%.tar}" + rm -f "${NAME}" + done + cd .. + cd .. +fi \ No newline at end of file diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..6b915132715aeff0fbc0cf547488d6db5d1cac36 --- /dev/null +++ b/src/imagenet/datasets.py @@ -0,0 +1,142 @@ +# Copyright (c) 2015-present, Facebook, Inc. +# All rights reserved. +import os +import json +import numpy as np + +from torchvision import datasets, transforms +from torchvision.datasets.folder import ImageFolder, default_loader +from torch.utils.data import Dataset, DataLoader + +# from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD +from timm.data import create_transform + + +PATH_TO_IMAGENET_VAL = '/sys/fs/cgroup/imagenet/val' + +def create_symlinks_to_imagenet(imagenet_folder, folder_to_scan): + if not os.path.exists(imagenet_folder): + os.makedirs(imagenet_folder) + folders_of_interest = os.listdir(folder_to_scan) + path_prefix = PATH_TO_IMAGENET_VAL + for folder in folders_of_interest: + os.symlink(path_prefix + folder, imagenet_folder+folder, target_is_directory=True) + +class INatDataset(ImageFolder): + def __init__(self, root, train=True, year=2018, transform=None, target_transform=None, + category='name', loader=default_loader): + self.transform = transform + self.loader = loader + self.target_transform = target_transform + self.year = year + # assert category in ['kingdom','phylum','class','order','supercategory','family','genus','name'] + path_json = os.path.join(root, f'{"train" if train else "val"}{year}.json') + with open(path_json) as json_file: + data = json.load(json_file) + with open(os.path.join(root, 'categories.json')) as json_file: + data_catg = json.load(json_file) + path_json_for_targeter = os.path.join(root, f"train{year}.json") + with open(path_json_for_targeter) as json_file: + data_for_targeter = json.load(json_file) + targeter = {} + indexer = 0 + for elem in data_for_targeter['annotations']: + king = [] + king.append(data_catg[int(elem['category_id'])][category]) + if king[0] not in targeter.keys(): + targeter[king[0]] = indexer + indexer += 1 + self.nb_classes = len(targeter) + self.samples = [] + for elem in data['images']: + cut = elem['file_name'].split('/') + target_current = int(cut[2]) + path_current = os.path.join(root, cut[0], cut[2], cut[3]) + categors = data_catg[target_current] + target_current_true = targeter[categors[category]] + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR10': + dataset = datasets.CIFAR10(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 10 + if args.data_set == 'CIFAR100': + dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': + if(is_generalization): root = args.data_path + else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': + dataset = INatDataset(args.data_path, train=is_train, year=2018, + category=args.inat_category, transform=transform) + nb_classes = dataset.nb_classes + elif args.data_set == 'INAT19': + dataset = INatDataset(args.data_path, train=is_train, year=2019, + category=args.inat_category, transform=transform) + nb_classes = dataset.nb_classes + return dataset, nb_classes + +def build_transform(is_train, args): + resize_im = args.input_size > 32 + if args.data_set == 'CIFAR100' or args.data_set == 'CIFAR10': + if is_train: + return transforms.Compose([ + transforms.RandomCrop(args.input_size, padding=4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), + ]) + else: + return transforms.Compose([ + transforms.RandomCrop(args.input_size, padding=4), + transforms.ToTensor(), + transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), + ]) + else: + if is_train: + # this should always dispatch to transforms_imagenet_train + transform = create_transform( + input_size=args.input_size, + is_training=True, + color_jitter=args.color_jitter, + auto_augment=args.aa, + interpolation=args.train_interpolation, + re_prob=args.reprob, + re_mode=args.remode, + re_count=args.recount, + ) + if not resize_im: + # replace RandomResizedCropAndInterpolation with + # RandomCrop + transform.transforms[0] = transforms.RandomCrop( + args.input_size, padding=4) + return transform + t = [] + if resize_im: + size = int((256 / 224) * args.input_size) + t.append( + transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images + ) + t.append(transforms.CenterCrop(args.input_size)) + t.append(transforms.ToTensor()) + t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) + return transforms.Compose(t) + + +class CIFARC_Dataset(Dataset): + def __init__(self, data_path, transform=None): + self.data = np.load(data_path) # (N, H, W, C) + self.transform = transform + + def __len__(self): + return self.data.shape[0] + + def __getitem__(self, idx): + img = self.data[idx] # (H, W, C) + if self.transform: + img = self.transform(img) + return img diff --git a/src/imagenet/finetune.py b/src/imagenet/finetune.py new file mode 100644 index 0000000000000000000000000000000000000000..d8d2cdcff7e9b082718455689ee89d1f033fb0d2 --- /dev/null +++ b/src/imagenet/finetune.py @@ -0,0 +1,325 @@ +import argparse +import os +import jax +import flax +import copy +import torch +import wandb +import optax +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from typing import Any, Dict, List +from flax.jax_utils import replicate, unreplicate +from flax.core.frozen_dict import freeze, unfreeze +from flax.training import checkpoints, train_state +from flax.traverse_util import flatten_dict, unflatten_dict +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from transformers.models.vit.modeling_flax_vit import ViTConfig +from datasets import build_dataset +import multiprocessing as mp +from pprint import pprint +import json +import shutil + +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +# ---------- Dataset Loader ---------- +def load_flax_params(checkpoint_dir, target): + msgpack_path = os.path.join(checkpoint_dir, "flax_model.msgpack") + with open(msgpack_path, "rb") as f: + packed_bytes = f.read() + params = from_bytes(target, packed_bytes) + return params +def parse_step_from_path(path: str) -> int: + m = re.search(r'last_(\d+)', path) + return int(m.group(1)) if m else 0 +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +# ---------- Training Utilities ---------- +def get_trainable_mask(params, config): + def is_att_param(keys): + if len(keys) < 5: return False + if keys[0] == "vit" and keys[1] == "encoder" and keys[2] == "layer" and int(keys[3]) in config.lmc_layer_indices and keys[4] == "attention": return True + if config.finetune_mlp == True and keys[0] == "vit" and keys[1] == "encoder" and keys[2] == "layer" and int(keys[3]) in config.lmc_layer_indices and keys[4] == "moe": return True + return False + def label_fn(path, _): + keys = [str(k.key) for k in path] + if is_att_param(keys): return "trainable" + return "frozen" + return jax.tree_util.tree_map_with_path(label_fn, params) +def pretrained2finetune_params(pretrained_params, finetune_params, config): + pretrained_params = unfreeze(pretrained_params) + finetune_params = unfreeze(finetune_params) + # 1. Copy top-level params (embeddings, layernorm, classifier) + finetune_params["vit"]["embeddings"] = copy.deepcopy(pretrained_params["vit"]["embeddings"]) + finetune_params["vit"]["layernorm"] = copy.deepcopy(pretrained_params["vit"]["layernorm"]) + finetune_params["classifier"] = copy.deepcopy(pretrained_params["classifier"]) + # 2. Copy encoder layers + for i in range(config.num_hidden_layers): + str_i = str(i) + if i in config.lmc_layer_indices: + ref_layer = pretrained_params["vit"]["encoder"]["layer"][str_i] # use layer 0 from pretrained + target_layer = finetune_params["vit"]["encoder"]["layer"][str_i] + # Copy shared parts + target_layer["layernorm_before"] = copy.deepcopy(ref_layer["layernorm_before"]) + target_layer["layernorm_after"] = copy.deepcopy(ref_layer["layernorm_after"]) + if config.finetune_mlp == False: + target_layer["moe"] = copy.deepcopy(ref_layer["moe"]) + # #Copy Attention weights + # target_layer["attention"] = copy.deepcopy(ref_layer["attention"]) + else: + finetune_params["vit"]["encoder"]["layer"][str_i] = copy.deepcopy(pretrained_params["vit"]["encoder"]["layer"][str_i]) + return freeze(finetune_params) + + +# ---------- Main ---------- +def main(args: argparse.Namespace): + if os.path.exists(args.model_path): + config = ViTConfig.from_pretrained(os.path.dirname(args.model_path)) + else: + raise FileNotFoundError(f"Config directory does not exist: {os.path.dirname(args.model_path)}") + os.makedirs(args.wandb_run_dir, exist_ok=True) + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name=f"finetune{args.lr}-{config.position_embeddings}-indice{','.join(str(i) for i in args.lmc_layer_indices)}-heads{args.num_attention_heads}" + f"-shared{config.num_shared_experts}-routed{config.num_routed_experts}-topk{config.topk}-mlp{str(args.finetune_mlp)}-seed{args.seed}", + save_code=True + ) + save_path = os.path.join( + args.save_dir,f"finetune{args.lr}-{config.position_embeddings}-indice{','.join(str(i) for i in args.lmc_layer_indices)}-heads{args.num_attention_heads}" + f"-shared{config.num_shared_experts}-routed{config.num_routed_experts}-topk{config.topk}-mlp{str(args.finetune_mlp)}-seed{args.seed}" + ) + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + pretrained_model = LMCFlaxViTForImageClassification(config, dtype=jnp.dtype(args.dtype)) + lmc_config = copy.deepcopy(config) + new_config = copy.deepcopy(config) + lmc_config.num_attention_heads = args.num_attention_heads + lmc_config.routed_scaling_factor = args.routed_scaling_factor + new_config.lmc_config = lmc_config + new_config.lmc_layer_indices = args.lmc_layer_indices + new_config.finetune_mlp = args.finetune_mlp + if os.path.exists(args.model_path): + pretrained_params = checkpoints.restore_checkpoint(ckpt_dir=args.model_path, target={"params": pretrained_model.params})["params"] + else: + raise FileNotFoundError(f"Checkpoint path not found: {args.model_path}") + pretrained_model.params = pretrained_params + # --- Initialize fine-tuning model --- + model = LMCFlaxViTForImageClassification(new_config,input_shape=(1,new_config.image_size, new_config.image_size, new_config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype),) + model.config.save_pretrained(save_path) + print_model(model.params) + model.params = pretrained2finetune_params(pretrained_model.params,model.params,new_config) + # model = pretrained_model + label_mask = get_trainable_mask(model.params,new_config) + print(json.dumps(label_mask, indent=2)) + num_train_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_train_steps, + end_value=args.min_lr, + ) + tx = optax.multi_transform( + transforms={ + 'trainable': optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay + ), + 'frozen': optax.set_to_zero() + }, + param_labels=label_mask + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + restored_params = load_flax_params(args.restore_checkpoint_path,state.params) + restored_step = parse_step_from_path(args.restore_checkpoint_path) + state = state.replace(params=restored_params, step=restored_step) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + train_metrics_stack = [] + train_loss = 0.0 + global_step = latest_global_step + best_val_acc1 = 0.0 + ###JUST FOR TESTING#### + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + print(f"Validation Loss: {val_loss:.4f}, Top-1 Accuracy: {val_acc1:.2f}%, Top-5 Accuracy: {val_acc5:.2f}%") + #### START FINETUNING #### + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + model.params = unreplicate(state).params + best_dir = os.path.join(save_path, f"best_{global_step}") + model.save_pretrained(best_dir) + print(f"Best model saved at step {global_step}") + remove_old_dirs_with_prefix(save_path, "best_", global_step) + # checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1,overwrite=True,orbax_checkpointer=None) + # Save last model + model.params = unreplicate(state).params + last_dir = os.path.join(save_path, f"last_{global_step}") + model.save_pretrained(last_dir) + print(f"Last model saved at step {global_step}") + remove_old_dirs_with_prefix(save_path, "last_", global_step) + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--model-path", type=str, default="", help="Path of Pretrained Model") + parser.add_argument("--num-attention-heads", type=int, default = 1) + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--finetune-mlp",action="store_true",help="Enable fine-tuning for the MLP. Default is False.") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=3, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR10','CIFAR100', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=8) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default="", help="wandb entity for logging") + parser.add_argument("--wandb-group", default="", help="wandb group for logging") + parser.add_argument("--wandb-project", default="", help="wandb project name for logging") + parser.add_argument("--wandb-run-dir", default=".wandb", help="wandb run dir") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) diff --git a/src/imagenet/generalization.py b/src/imagenet/generalization.py new file mode 100644 index 0000000000000000000000000000000000000000..75b42f3e1f9a64eb91d6d6deeb13ab70278a0e2e --- /dev/null +++ b/src/imagenet/generalization.py @@ -0,0 +1,212 @@ +import argparse +import os +import jax +import flax +import copy +import torch +import wandb +import optax +import itertools +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from jax import random, vmap +from flax import linen as nn +from typing import Any, Dict, List +from jax import random, vmap +from flax.jax_utils import replicate, unreplicate +from flax.core.frozen_dict import freeze, unfreeze +from flax.training import checkpoints, train_state +from flax.traverse_util import flatten_dict, unflatten_dict +from flax.training.common_utils import get_metrics, onehot, shard +from utils import flatten_params, lerp, unflatten_params +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from transformers.models.vit.modeling_flax_vit import ViTConfig +from matching_utils import weight_matching_attn +from datasets import build_dataset +from flax.serialization import from_bytes +import multiprocessing as mp +import matplotlib.pyplot as plt +import numpy as np +import json +mp.set_start_method("spawn", force=True) +def load_flax_params(checkpoint_dir, target): + msgpack_path = os.path.join(checkpoint_dir, "flax_model.msgpack") + with open(msgpack_path, "rb") as f: + packed_bytes = f.read() + params = from_bytes(target, packed_bytes) + return params + +def imagenet_data_loader(args,is_generalization): + dataset_val, args.nb_classes = build_dataset(is_train=False, args=args,is_generalization=is_generalization) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_val +def sample_random_batch(dataloader, seed=0): + """Return (images, labels) from a random batch without materializing the loader.""" + n = len(dataloader) # DataLoader should define __len__ + idx = int(np.random.default_rng(seed).integers(n)) + return next(itertools.islice(iter(dataloader), idx, None)) +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def prepare_image_batch(images: torch.Tensor, labels: torch.Tensor): + # torch -> numpy + images = images.detach().cpu().numpy() # (B, C, H, W) + labels = labels.detach().cpu().numpy() + images = np.transpose(images, (0, 2, 3, 1)) # (B, H, W, C) + images = images.astype(np.float32) + labels = labels.astype(np.int32) + return {"images": shard(images), "labels": shard(labels)} +def make_stuff(model): + apply_fn = model.__call__ + def batch_eval(params, batch): + outputs = apply_fn(params=params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_batch_eval = jax.pmap(batch_eval, axis_name="batch") + def dataset_loss_and_acc(params, dataloader): + """Evaluate over the dataloader on all devices""" + eval_results = [] + pbar = tqdm(dataloader, desc="Evaluating", leave=False) + rep_params = replicate(params) + for step_idx, (images, labels) in enumerate(pbar): + batch = prepare_image_batch(images, labels) + eval_metric = parallel_batch_eval(rep_params, batch) + eval_results.append(copy.deepcopy(eval_metric)) + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + loss, acc1 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]) + return loss, acc1 + def get_attention_inputs(params, dataloader): + activation = {} + for images, labels in dataloader: + images_host = jnp.array(images) + _ = apply_fn(params=params,pixel_values=images_host,train=False,attention_input=activation,) + break + return activation + return {"batch_eval": batch_eval, "dataset_loss_and_acc": dataset_loss_and_acc, "get_attention_inputs": get_attention_inputs} +def compute_interpolation(params_a, params_b_target, lambdas, stuff, test_loader, desc="Interpolation"): + train_loss_interp, test_loss_interp = [], [] + train_acc_interp, test_acc_interp = [], [] + for lam in tqdm(lambdas, desc=desc): + p_interp = freeze(lerp(lam, unfreeze(params_a), unfreeze(params_b_target))) + test_loss, test_acc = stuff["dataset_loss_and_acc"](p_interp, test_loader) + test_loss_interp.append(test_loss) + test_acc_interp.append(test_acc) + return {"Test Loss": [float(f"{x:.4f}") for x in test_loss_interp],"Test Acc": [float(f"{x:.4f}") for x in test_acc_interp]} +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-a", type=str, required=True, help="Path to first fine-tuned GPT2 model checkpoint") + parser.add_argument("--model-b", type=str, required=True, help="Path to second fine-tuned GPT@ model checkpoint") + parser.add_argument("--batch-size", type=int, default=256) + parser.add_argument("--data-path", type=str, default=None) + parser.add_argument("--data-original", type=str, required=True) + parser.add_argument("--data-generalization", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=8) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--save-path", type=str, default="/", help="Path to plot directory") + args = parser.parse_args() + config = ViTConfig.from_json_file(os.path.join(os.path.dirname(args.model_a).rstrip("/"),'config.json')) + lmc_config = ViTConfig(**config.lmc_config) + config.lmc_config = lmc_config + model = LMCFlaxViTForImageClassification(config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype),) + print_model(model.params) + if os.path.exists(args.model_a) and os.path.exists(args.model_b): + params_a = load_flax_params(args.model_a,copy.deepcopy(model.params)) + params_b = load_flax_params(args.model_b,copy.deepcopy(model.params)) + else: + raise FileNotFoundError(f"Checkpoint path does not exist") + stuff = make_stuff(model = model) + args.data_path = args.data_original + original_val_loader = imagenet_data_loader(args,False) + args.data_path = args.data_generalization + generalization_val_loader = imagenet_data_loader(args,True) + lambdas = jnp.linspace(0, 1, num=3) + rng = random.PRNGKey(args.seed) + # Compute naive interpolation + # all_results = {"Naive": naive_results} + naive_original_results = compute_interpolation(params_a, params_b, lambdas, stuff, original_val_loader, desc="Naive Original Interpolation") + print(json.dumps({"Naive-Original": naive_original_results}, indent=2)) + naive_original_results = compute_interpolation(params_a, params_b, lambdas, stuff, generalization_val_loader, desc="Naive Generalization Interpolation") + print(json.dumps({"Naive-Generalization": naive_original_results}, indent=2)) + exit() + # Compute weight matching interpolations for each method + activation = None #stuff["get_attention_inputs"](params_b, val_loader) + aligned_models = weight_matching_attn(rng, params_a, params_b, activation, config) + for method, params_b_aligned in aligned_models.items(): + original_method_results = compute_interpolation(params_a, params_b_aligned, lambdas, stuff, original_val_loader, desc=f"{method} Interpolation") + print(json.dumps({method+"original": original_method_results}, indent=2)) + generalization_method_results = compute_interpolation(params_a, params_b_aligned, lambdas, stuff, generalization_val_loader, desc=f"{method} Interpolation") + print(json.dumps({method+"generalization": generalization_method_results}, indent=2)) + exit() + # Save directories + os.makedirs(f"./plots/imagenet", exist_ok=True) + os.makedirs(f"./results/imagenet", exist_ok=True) + # Save results JSON + print("Save List of Values...") + name_a = os.path.basename(os.path.dirname(args.model_a).rstrip("/")) + name_b = os.path.basename(os.path.dirname(args.model_b).rstrip("/")) + result_path = f'results/imagenet/{name_a}+{name_b}.json' + with open(result_path, 'w') as f: + json.dump(all_results, f, indent=2) + # Plot + print("Generating plots...") + plot_path = f"./plots/imagenet/{name_a}+{name_b}.pdf" + plt.rcParams.update({ + "font.family": "serif", + 'legend.frameon': False, + 'lines.linewidth': 2, + 'font.size': 13, + 'axes.labelsize': 16, + 'xtick.labelsize': 11, + 'ytick.labelsize': 11, + 'legend.fontsize': 11, + }) + plt.style.use('tableau-colorblind10') + num_points = len(all_results["Naive"]["Test Loss"]) + lambda_values = np.linspace(0, 1, num_points) + fig, axs = plt.subplots(2, 2, figsize=(12, 10)) + metrics = ["Test Loss", "Test Acc", "Test Loss", "Test Acc",] + positions = [(0, 0), (0, 1), (1, 0), (1, 1)] + for metric, pos in zip(metrics, positions): + row, col = pos + ax = axs[row, col] + for method in all_results: + ax.plot(lambda_values, all_results[method][metric], label=method) + ax.set_xticks([0, 0.5, 1]) + ax.set_xticklabels(["Model 1", r"$\lambda$", "Model 2"]) + ax.set_ylabel(metric) + ax.legend(loc='best') + plt.tight_layout() + plt.savefig(plot_path) + plt.close() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/imagenet/lmc_model.py b/src/imagenet/lmc_model.py new file mode 100644 index 0000000000000000000000000000000000000000..7df39ef50d3c94e93922e5117f44abbcf7459708 --- /dev/null +++ b/src/imagenet/lmc_model.py @@ -0,0 +1,595 @@ +from typing import Optional, Tuple, List +import flax.linen as nn +import jax +import copy +import numpy as np +import jax.numpy as jnp +from flax.linen.attention import dot_product_attention_weights +from flax.core.frozen_dict import FrozenDict, freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ( + FlaxViTPreTrainedModel, + ViTConfig, + FlaxBaseModelOutput, + FlaxBaseModelOutputWithPooling, + FlaxViTPatchEmbeddings, + FlaxViTPooler, + ACT2FN, + FlaxViTIntermediate, + FlaxViTOutput, + FlaxPreTrainedModel, + FlaxSequenceClassifierOutput, +) +from typing import Callable + +def print_model(flax_params, file=None): + flat_params = flatten_dict(flax_params) + for path, value in flat_params.items(): + name = "/".join(path) + line = f"{name} {value.shape}" + if file: + print(line, file=file) + else: + print(line) +def print_model_with_prefix(flax_params, prefix: str, file=None): + flat_params = flatten_dict(flax_params) + for path, value in flat_params.items(): + name = ".".join(path) + if name.startswith(prefix): + line = f"{name} {value.shape} \n {value} \n \n" + if file: print(line, file=file) + else: print(line) +def create_sinusoidal_positions(n_pos, dim): + position_enc = np.array([[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)]) + sentinel = dim // 2 + dim % 2 + out = np.zeros_like(position_enc) + out[:, 0:sentinel] = np.sin(position_enc[:, 0::2]) + out[:, sentinel:] = np.cos(position_enc[:, 1::2]) + + return jnp.array(out) + +class LMCFlaxViTEmbeddings(nn.Module): + """Construct the CLS token, position and patch embeddings.""" + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + self.cls_token = self.param( + "cls_token", + jax.nn.initializers.variance_scaling(self.config.initializer_range**2, "fan_in", "truncated_normal"), + (1, 1, self.config.hidden_size), + ) + self.patch_embeddings = FlaxViTPatchEmbeddings(self.config, dtype=self.dtype) + num_patches = self.patch_embeddings.num_patches + if self.config.position_embeddings == "learnable": + self.position_embeddings = self.param( + "position_embeddings", + jax.nn.initializers.variance_scaling(self.config.initializer_range**2, "fan_in", "truncated_normal"), + (1, num_patches + 1, self.config.hidden_size), + ) + elif self.config.position_embeddings == "sinusoidal": + self.position_embeddings = jnp.expand_dims(create_sinusoidal_positions(num_patches + 1, self.config.hidden_size),axis=0) + self.dropout = nn.Dropout(rate=self.config.hidden_dropout_prob) + def __call__(self, pixel_values, deterministic=True): + batch_size = pixel_values.shape[0] + embeddings = self.patch_embeddings(pixel_values) + cls_tokens = jnp.broadcast_to(self.cls_token, (batch_size, 1, self.config.hidden_size)) + embeddings = jnp.concatenate((cls_tokens, embeddings), axis=1) + if self.config.position_embeddings in ["learnable","sinusoidal"]: + embeddings = embeddings + self.position_embeddings + embeddings = self.dropout(embeddings, deterministic=deterministic) + return embeddings + +class LMCFlaxViTSelfAttention(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + if self.config.hidden_size % self.config.num_attention_heads != 0: + raise ValueError( + "`config.hidden_size`: {self.config.hidden_size} has to be a multiple of `config.num_attention_heads`:" + " {self.config.num_attention_heads}" + ) + self.query = nn.Dense( + self.config.hidden_size, + dtype=self.dtype, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, mode="fan_in", distribution="truncated_normal" + ), + use_bias=self.config.qkv_bias, + ) + self.key = nn.Dense( + self.config.hidden_size, + dtype=self.dtype, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, mode="fan_in", distribution="truncated_normal" + ), + use_bias=self.config.qkv_bias, + ) + self.value = nn.Dense( + self.config.hidden_size, + dtype=self.dtype, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, mode="fan_in", distribution="truncated_normal" + ), + use_bias=self.config.qkv_bias, + ) + if self.config.position_embeddings == "rope": + self.rotary_value = self.config.rotary_value + def __call__(self, hidden_states, sinusoidal_pos, deterministic: bool = True, output_attentions: bool = False): + head_dim = self.config.hidden_size // self.config.num_attention_heads + query_states = self.query(hidden_states).reshape( + hidden_states.shape[:2] + (self.config.num_attention_heads, head_dim) + ) + value_states = self.value(hidden_states).reshape( + hidden_states.shape[:2] + (self.config.num_attention_heads, head_dim) + ) + key_states = self.key(hidden_states).reshape( + hidden_states.shape[:2] + (self.config.num_attention_heads, head_dim) + ) + if sinusoidal_pos is not None and self.config.position_embeddings == 'rope': + apply_sinusoidal_pos = create_sinusoidal_positions(sinusoidal_pos.shape[0], head_dim) + if self.rotary_value: + query_states, key_states, value_states = self.apply_rotary_position_embeddings( + apply_sinusoidal_pos, query_states, key_states, value_states + ) + else: + query_states, key_states = self.apply_rotary_position_embeddings( + apply_sinusoidal_pos, query_states, key_states + ) + dropout_rng = None + if not deterministic and self.config.attention_probs_dropout_prob > 0.0: + dropout_rng = self.make_rng("dropout") + attn_weights = dot_product_attention_weights( + query_states, + key_states, + dropout_rng=dropout_rng, + dropout_rate=self.config.attention_probs_dropout_prob, + broadcast_dropout=True, + deterministic=deterministic, + dtype=self.dtype, + precision=None, + ) + attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value_states) + attn_output = attn_output.reshape(attn_output.shape[:2] + (-1,)) + outputs = (attn_output, attn_weights) if output_attentions else (attn_output,) + return outputs + @staticmethod + def apply_rotary_position_embeddings(sinusoidal_pos, query_layer, key_layer, value_layer=None): + sin, cos = jnp.split(sinusoidal_pos, 2, axis=-1) + sin_pos = jnp.stack([sin, sin], axis=-1).reshape(sinusoidal_pos.shape) + cos_pos = jnp.stack([cos, cos], axis=-1).reshape(sinusoidal_pos.shape) + def rotate_layer(layer, sin_pos, cos_pos): + rotate_half_layer = jnp.stack([-layer[..., 1::2], layer[..., ::2]], axis=-1).reshape(layer.shape) + rotary_matrix_cos = jnp.einsum("bslh,...sh->bslh", layer, cos_pos) + rotary_matrix_sin = jnp.einsum("bslh,...sh->bslh", rotate_half_layer, sin_pos) + return rotary_matrix_cos + rotary_matrix_sin + query_layer = rotate_layer(query_layer, sin_pos, cos_pos) + key_layer = rotate_layer(key_layer, sin_pos, cos_pos) + if value_layer is not None: + value_layer = rotate_layer(value_layer, sin_pos, cos_pos) + return query_layer, key_layer, value_layer + return query_layer, key_layer +class LMCFlaxViTSelfOutput(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + self.dense = nn.Dense( + self.config.hidden_size, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, "fan_in", "truncated_normal" + ), + dtype=self.dtype, + ) + self.dropout = nn.Dropout(rate=self.config.hidden_dropout_prob) + + def __call__(self, hidden_states, input_tensor, deterministic: bool = True): + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states, deterministic=deterministic) + return hidden_states +class LMCFlaxViTAttention(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.attention = LMCFlaxViTSelfAttention(self.config, dtype=self.dtype) + self.output = LMCFlaxViTSelfOutput(self.config, dtype=self.dtype) + def __call__(self, hidden_states, sinusoidal_pos, deterministic=True, output_attentions: bool = False): + attn_outputs = self.attention(hidden_states, sinusoidal_pos, deterministic=deterministic, output_attentions=output_attentions) + attn_output = attn_outputs[0] + hidden_states = self.output(attn_output, hidden_states, deterministic=deterministic) + outputs = (hidden_states,) + if output_attentions: + outputs += (attn_outputs[1],) + + return outputs + + +class LMCFlaxViTMLP(nn.Module): + config: ViTConfig + intermediate_size : int + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + self.intermediate = nn.Dense( + self.intermediate_size, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, "fan_in", "truncated_normal" + ), + dtype=self.dtype, + ) + self.activation = ACT2FN[self.config.hidden_act] + self.output = nn.Dense( + self.config.hidden_size, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, "fan_in", "truncated_normal" + ), + dtype=self.dtype, + ) + self.dropout = nn.Dropout(rate=self.config.hidden_dropout_prob) + def __call__(self, layer_output, deterministic: bool = True): + hidden_states = self.intermediate(layer_output) + hidden_states = self.activation(hidden_states) + hidden_states = self.output(hidden_states) + hidden_states = self.dropout(hidden_states, deterministic=deterministic) + return hidden_states + +class LMCFlaxViTRouter(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 + def setup(self): + self.n_routed_experts = self.config.num_routed_experts + self.n_group = getattr(self.config, "n_group", 1) + self.topk_group = getattr(self.config, "topk_group", 1) + self.top_k = self.config.topk + self.routed_scaling_factor = self.config.routed_scaling_factor + self.norm_topk_prob = getattr(self.config, "norm_topk_prob", False) + # Weight and bias for router computation + kernel_init = jax.nn.initializers.normal(self.config.initializer_range) + self.router_weight = self.param( + "router_weight", kernel_init, (self.n_routed_experts, self.config.hidden_size) + ) + self.router_bias = self.param( + "router_bias", lambda rng, shape: jnp.zeros(shape, dtype=self.dtype), (self.n_routed_experts,) + ) + self.e_score_correction_bias = self.param( + "e_score_correction_bias", lambda rng, shape: jnp.zeros(shape, dtype=self.dtype), (self.n_routed_experts,) + ) + + def get_topk_indices(self, scores): + scores_for_choice = scores + self.e_score_correction_bias[None, :] + scores_grouped = scores_for_choice.reshape(-1, self.n_group, self.n_routed_experts // self.n_group) + top2_scores = jax.lax.top_k(scores_grouped, 2)[0] + group_scores = jnp.sum(top2_scores, axis=-1) + top_group_scores, group_idx = jax.lax.top_k(group_scores, self.topk_group) + group_mask = jnp.zeros_like(group_scores) + group_mask = group_mask.at[jnp.arange(group_mask.shape[0])[:, None], group_idx].set(1) + group_mask_expanded = jnp.repeat(group_mask[:, :, None], self.n_routed_experts // self.n_group, axis=-1) + score_mask = group_mask_expanded.reshape(-1, self.n_routed_experts) + scores_for_choice = jnp.where( + score_mask, + scores_for_choice, + jnp.zeros_like(scores_for_choice), + ) + topk_weights, topk_indices = jax.lax.top_k(scores_for_choice, self.top_k) + return topk_indices, topk_weights + + def __call__(self, hidden_states): + router_logits = jnp.matmul(hidden_states, self.router_weight.T) + self.router_bias + scores = jax.nn.sigmoid(router_logits) + if self.n_group > 1: + topk_indices, topk_weights = self.get_topk_indices(scores) + else: + topk_weights, topk_indices = jax.lax.top_k(scores, self.top_k) + if self.norm_topk_prob: + denominator = jnp.sum(topk_weights, axis=-1, keepdims=True) + 1e-20 + topk_weights = topk_weights / denominator + topk_weights = topk_weights * self.routed_scaling_factor + return topk_indices, topk_weights + +class LMCFlaxViTMoE(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + + def setup(self): + self.n_experts = self.config.num_routed_experts + self.n_shared_experts = self.config.num_shared_experts + self.top_k = self.config.topk + self.router = LMCFlaxViTRouter(config=self.config, dtype=self.dtype) # Using ViT router here + self.experts = [ + LMCFlaxViTMLP(config=self.config, intermediate_size=self.config.intermediate_size, dtype=self.dtype,) + for _ in range(self.n_experts) + ] + if self.n_shared_experts > 0: + self.shared_experts = LMCFlaxViTMLP(config=self.config, intermediate_size=self.config.intermediate_size * self.n_shared_experts, dtype=self.dtype) + def __call__(self, hidden_states, deterministic: bool = True): + residual = hidden_states + orig_shape = hidden_states.shape + if self.n_experts == 0 or self.top_k == 0: + hidden_states = self.shared_experts(hidden_states, deterministic=deterministic) + return hidden_states + hidden_states_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + topk_indices, topk_weights = self.router(hidden_states_flat) + expert_outputs = [] + for expert in self.experts: + output = expert(hidden_states_flat, deterministic=deterministic) # [n_tokens, hidden_dim] + expert_outputs.append(output) + expert_outputs = jnp.stack(expert_outputs, axis=1) # [n_tokens, n_experts, hidden_dim] + # Build routing mask: + routing_mask = jax.nn.one_hot(topk_indices, self.n_experts, dtype=self.dtype) + routing_mask = routing_mask.sum(axis=1) > 0 # [n_tokens, n_experts] + # Compute weights per expert + weights_per_expert = jax.nn.one_hot(topk_indices, self.n_experts, dtype=self.dtype) + weights_per_expert = (weights_per_expert * topk_weights[..., None]).sum(axis=1) + # Mask out experts not routed + weights_per_expert = weights_per_expert * routing_mask.astype(self.dtype) + # Multiply and sum + weighted_expert_outputs = expert_outputs * weights_per_expert[..., None] # [n_tokens, n_experts, hidden_dim] + final_output = weighted_expert_outputs.sum(axis=1) # [n_tokens, hidden_dim] + # Add shared expert + if self.n_shared_experts > 0: + shared_output = self.shared_experts(residual, deterministic=deterministic) + final_output = final_output.reshape(orig_shape) + shared_output + else: + final_output = final_output.reshape(orig_shape) + return final_output +class LMCFlaxViTLayer(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + self.attention = LMCFlaxViTAttention(self.config, dtype=self.dtype) + self.moe = LMCFlaxViTMoE(self.config, dtype=self.dtype) + self.layernorm_before = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype) + self.layernorm_after = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype) + def __call__(self, hidden_states,sinusoidal_pos, deterministic: bool = True, output_attentions: bool = False): + ln_hidden = self.layernorm_before(hidden_states) + attention_outputs = self.attention( + ln_hidden, sinusoidal_pos, deterministic=deterministic, output_attentions=output_attentions, + ) + attention_output = attention_outputs[0] + # first residual connection + attention_output = attention_output + hidden_states + # in ViT, layernorm is also applied after self-attention + layer_output = self.layernorm_after(attention_output) + hidden_states = self.moe(layer_output, deterministic=deterministic) + hidden_states = hidden_states + attention_output + outputs = (hidden_states,) + if output_attentions: + outputs += (attention_outputs[1],) + return outputs + +class LMCFlaxViTLayerCollection(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + self.layers = [ + LMCFlaxViTLayer(self.config.lmc_config, name=str(i), dtype=self.dtype) if i in self.config.lmc_layer_indices + else LMCFlaxViTLayer(self.config, name=str(i), dtype=self.dtype) + for i in range(self.config.num_hidden_layers) + ] + def __call__( + self, + hidden_states, + sinusoidal_pos, + deterministic: bool = True, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ): + all_attentions = () if output_attentions else None + all_hidden_states = () if output_hidden_states else None + + for i, layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = layer(hidden_states,sinusoidal_pos,deterministic=deterministic, output_attentions=output_attentions) + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions += (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + outputs = (hidden_states,) + if not return_dict: + return tuple(v for v in outputs if v is not None) + + return FlaxBaseModelOutput( + last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions + ) +class LMCFlaxViTEncoder(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + num_patches = (self.config.image_size**2)//self.config.patch_size + self.embed_positions = create_sinusoidal_positions(num_patches + 1, self.config.hidden_size//self.config.num_attention_heads) + self.layer = LMCFlaxViTLayerCollection(self.config, dtype=self.dtype) + def __call__( + self, + hidden_states, + deterministic: bool = True, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ): + sinusoidal_pos = self.embed_positions[: hidden_states.shape[1], :] + return self.layer( + hidden_states, + sinusoidal_pos, + deterministic=deterministic, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + +class LMCFlaxViTPreTrainedModel(FlaxPreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + config_class = ViTConfig + base_model_prefix = "vit" + main_input_name = "pixel_values" + module_class: nn.Module = None + def __init__( + self, + config: ViTConfig, + input_shape=None, + seed: int = 0, + dtype: jnp.dtype = jnp.float32, + _do_init: bool = True, + **kwargs, + ): + module = self.module_class(config=config, dtype=dtype, **kwargs) + if input_shape is None: + input_shape = (1, config.image_size, config.image_size, config.num_channels) + super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init) + + def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict: + # init input tensors + pixel_values = jnp.zeros(input_shape, dtype=self.dtype) + + params_rng, dropout_rng = jax.random.split(rng) + rngs = {"params": params_rng, "dropout": dropout_rng} + + random_params = self.module.init(rngs, pixel_values, return_dict=False)["params"] + + if params is not None: + random_params = flatten_dict(unfreeze(random_params)) + params = flatten_dict(unfreeze(params)) + for missing_key in self._missing_keys: + params[missing_key] = random_params[missing_key] + self._missing_keys = set() + return freeze(unflatten_dict(params)) + else: + return random_params + + def __call__( + self, + pixel_values, + params: Optional[dict] = None, + dropout_rng: jax.random.PRNGKey = None, + train: bool = False, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ): + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + pixel_values = jnp.transpose(pixel_values, (0, 2, 3, 1)) + # Handle any PRNG if needed + rngs = {} + if dropout_rng is not None: + rngs["dropout"] = dropout_rng + + return self.module.apply( + {"params": params or self.params}, + jnp.array(pixel_values, dtype=jnp.float32), + not train, + output_attentions, + output_hidden_states, + return_dict, + rngs=rngs, + ) + + +class LMCFlaxViTModule(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + add_pooling_layer: bool = True + + def setup(self): + self.embeddings = LMCFlaxViTEmbeddings(self.config, dtype=self.dtype) + self.encoder = LMCFlaxViTEncoder(self.config, dtype=self.dtype) + self.layernorm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype) + self.pooler = FlaxViTPooler(self.config, dtype=self.dtype) if self.add_pooling_layer else None + + def __call__( + self, + pixel_values, + deterministic: bool = True, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ): + hidden_states = self.embeddings(pixel_values, deterministic=deterministic) + + outputs = self.encoder( + hidden_states, + deterministic=deterministic, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = outputs[0] + hidden_states = self.layernorm(hidden_states) + pooled = self.pooler(hidden_states) if self.add_pooling_layer else None + + if not return_dict: + # if pooled is None, don't return it + if pooled is None: + return (hidden_states,) + outputs[1:] + return (hidden_states, pooled) + outputs[1:] + + return FlaxBaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=pooled, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + + +class LMCFlaxViTModel(LMCFlaxViTPreTrainedModel): + module_class = LMCFlaxViTModule + +class LMCFlaxViTForImageClassificationModule(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 + def setup(self): + self.vit = LMCFlaxViTModule(config=self.config, dtype=self.dtype, add_pooling_layer=False) + self.classifier = nn.Dense( + self.config.num_labels, + dtype=self.dtype, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, "fan_in", "truncated_normal" + ), + ) + + def __call__( + self, + pixel_values=None, + deterministic: bool = True, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.vit( + pixel_values, + deterministic=deterministic, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + logits = self.classifier(hidden_states[:, 0, :]) + + if not return_dict: + output = (logits,) + outputs[2:] + return output + + return FlaxSequenceClassifierOutput( + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) +class LMCFlaxViTForImageClassification(LMCFlaxViTPreTrainedModel): + module_class = LMCFlaxViTForImageClassificationModule + diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4f47b19aac0cd1db8c6a8f5217da6f127751fb02 --- /dev/null +++ b/src/imagenet/matching_utils.py @@ -0,0 +1,484 @@ +import os +import jax +import time +import copy +import jax.nn as nn +import jax.lax as lax +import jax.numpy as jnp +from lmc_model import print_model +from collections import defaultdict +from typing import NamedTuple +from flax.core import freeze, unfreeze +from jax import random, tree_util, jit, grad, value_and_grad +from scipy.optimize import linear_sum_assignment, minimize +import numpy as np +import matplotlib.pyplot as plt +def compute_objective(A, X, X_prime, Y, Y_prime): + A_inv = np.linalg.inv(A) + term1 = X - X_prime @ A.T + term2 = Y - Y_prime @ A_inv + return np.sum(term1**2) + np.sum(term2**2) + +def compute_gradient(A, X, X_prime, Y, Y_prime): + A_inv = np.linalg.inv(A) + term1 = -2 * X.T @ X_prime + 2 * A @ X_prime.T @ X_prime + term2 = 2 * A_inv.T @ Y_prime.T @ (Y - Y_prime @ A_inv) @ A_inv.T + return term1 + term2 + +def line_search(A, grad, X, X_prime, Y, Y_prime, max_step=1, tau=0.5, c1=1e-4): + eta = max_step + f_current = compute_objective(A, X, X_prime, Y, Y_prime) + grad_norm2 = np.sum(grad**2) + n = A.shape[0] + while eta > 1e-10: + A_new = A - eta * grad + if np.linalg.matrix_rank(A_new) < n: + eta *= tau + continue + f_new = compute_objective(A_new, X, X_prime, Y, Y_prime) + if f_new <= f_current - c1 * eta * grad_norm2: + return eta + eta *= tau + return 0 + +@jax.jit +def compute_objective_jax(A, X, X_prime, Y, Y_prime, cond_threshold=1e6): + cond = jnp.linalg.cond(A) + def safe_obj(): + A_inv = jnp.linalg.inv(A) + term1 = X - X_prime @ A.T + term2 = Y - Y_prime @ A_inv + return jnp.sum(term1**2) + jnp.sum(term2**2) + return lax.cond(cond > cond_threshold, lambda: jnp.inf, safe_obj) +compute_value_and_grad_jax = jit(value_and_grad(compute_objective_jax)) +def solve_orthogonal(X, X_prime, Y, Y_prime): + B = X.T @ X_prime + Y.T @ Y_prime + U, _, Vt = np.linalg.svd(B) + return U @ Vt + + +def optimize_alignment(A_init, X, X_prime, Y, Y_prime, max_iter=5000): + objective_values = [] + grad_norms = [] + condition_nums = [] + + def obj_fn(flat_A): + A = flat_A.reshape(A_init.shape) + obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime)) + return float(obj), np.array(grad_val).flatten() + + def callback(flat_A): + A = flat_A.reshape(A_init.shape) + obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime)) + grad_norm = jnp.linalg.norm(grad_val, 'fro') + cond = jnp.linalg.cond(jnp.array(A)) + objective_values.append(float(obj)) + grad_norms.append(float(grad_norm)) + condition_nums.append(float(cond)) + + res = minimize(obj_fn, A_init.flatten(), jac=True, method='L-BFGS-B', options={'maxiter': max_iter}, callback=callback) + A_opt = res.x.reshape(A_init.shape) + return A_opt, objective_values, grad_norms, condition_nums + +def extract_attention_params(attn): + key =np.array(attn['attention']['key']['kernel']) + key_bias = np.array(attn['attention']['key']['bias']) + query = np.array(attn['attention']['query']['kernel']) + query_bias = np.array(attn['attention']['query']['bias']) + value = np.array(attn['attention']['value']['kernel']) + value_bias = np.array(attn['attention']['value']['bias']) + out = np.array(attn['output']['dense']['kernel']) + out_bias = np.array(attn['output']['dense']['bias']) + return query, key, value, query_bias, key_bias, value_bias, out, out_bias +def reshape_attention_weights(query, key, value, query_bias, key_bias, value_bias, out_kernel, num_heads): + D = query.shape[0] + D_k = D_v = D // num_heads + def stack_per_head(tensor, axis=0): + return np.stack([ + tensor[:,i * D_k:(i + 1) * D_k] if axis == 0 else tensor[i * D_k:(i + 1) * D_k,:] + for i in range(num_heads) + ]) + def stack_bias_per_head(bias): + return np.stack([bias[i * D_k:(i + 1) * D_k] for i in range(num_heads)]) + W_Q = stack_per_head(query) + W_K = stack_per_head(key) + W_V = stack_per_head(value) + W_O = stack_per_head(out_kernel, axis=1) + b_Q = stack_bias_per_head(query_bias) + b_K = stack_bias_per_head(key_bias) + b_V = stack_bias_per_head(value_bias) + return W_Q, b_Q, W_K, b_K, W_V, b_V, W_O + +def compute_extended_weights(W, b): + return np.vstack([W, b.reshape(1, -1)]) + +def compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + h, activations, alpha=0.5): + C = np.zeros((h, h)) + for i in range(h): + tilde_W_Q_a_i = np.vstack([W_Q_a[i], b_Q_a[i].reshape(1, -1)]) + tilde_W_K_a_i = np.vstack([W_K_a[i], b_K_a[i].reshape(1, -1)]) + tilde_W_V_a_i = np.vstack([W_V_a[i], b_V_a[i].reshape(1, -1)]) + QKT_a_i = tilde_W_Q_a_i @ tilde_W_K_a_i.T + VO_a_i = tilde_W_V_a_i @ W_O_a[i] + centered_QKT_a_i = QKT_a_i - np.mean(QKT_a_i, axis=1, keepdims=True) + for j in range(h): + tilde_W_Q_b_j = np.vstack([W_Q_b[j], b_Q_b[j].reshape(1, -1)]) + tilde_W_K_b_j = np.vstack([W_K_b[j], b_K_b[j].reshape(1, -1)]) + tilde_W_V_b_j = np.vstack([W_V_b[j], b_V_b[j].reshape(1, -1)]) + QKT_b_j = tilde_W_Q_b_j @ tilde_W_K_b_j.T + VO_b_j = tilde_W_V_b_j @ W_O_b[j] + centered_QKT_b_j = QKT_b_j - np.mean(QKT_b_j, axis=1, keepdims=True) + cost = 0.5 * np.sum((centered_QKT_a_i - centered_QKT_b_j) ** 2) + cost += 0.5 * np.sum((VO_a_i - VO_b_j) ** 2) + C[i, j] = cost + return C +# def compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, +# W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, +# num_heads, activations, alpha=0.5): +# # activations: (B, L, D) +# B, L, D = activations.shape +# print("Activations shape: ", activations.shape) +# # Augment activations with a column of ones for bias calculation +# ones_col = jnp.ones((B, L, 1)) +# X_tilde = jnp.concatenate([activations, ones_col], axis=-1) # Shape (B, L, D+1) +# d_head = W_Q_a[0].shape[1] +# sqrt_d = jnp.sqrt(float(d_head)) +# C = np.zeros((num_heads, num_heads)) +# for i in range(num_heads): +# # Pre-compute for model A, head i +# tilde_W_Q_a_i = compute_extended_weights(W_Q_a[i], b_Q_a[i]) +# tilde_W_K_a_i = compute_extended_weights(W_K_a[i], b_K_a[i]) +# tilde_W_V_a_i = compute_extended_weights(W_V_a[i], b_V_a[i]) +# Q_a_i = X_tilde @ tilde_W_Q_a_i +# K_a_i = X_tilde @ tilde_W_K_a_i +# S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d +# P_a_i = nn.softmax(S_a_i, axis=-1) +# V_tilde_a_i = X_tilde @ tilde_W_V_a_i +# V_a_i = V_tilde_a_i @ W_O_a[i] +# for j in range(num_heads): +# # Compute for model B, head j +# tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j]) +# tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j]) +# tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j]) +# Q_b_j = X_tilde @ tilde_W_Q_b_j +# K_b_j = X_tilde @ tilde_W_K_b_j +# S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d +# P_b_j = nn.softmax(S_b_j, axis=-1) +# V_tilde_b_j = X_tilde @ tilde_W_V_b_j +# V_b_j = V_tilde_b_j @ W_O_b[j] +# cost_P = jnp.sum((P_a_i - P_b_j)**2) +# cost_V = jnp.sum((V_a_i - V_b_j)**2) +# total_cost = (alpha * cost_P + (1 - alpha) * cost_V) / B +# C[i, j] = total_cost +# return C +# def compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, +# W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, +# num_heads, activations, alpha=0.5, epsilon=1e-8): +# """ +# Computes the cost matrix for attention head permutation using Cosine Similarity. +# The cost C[i, j] is defined as 1 - cosine_similarity, which is minimized when +# the output tensors are most similar. +# """ +# # activations: (B, L, D) +# B, L, D = activations.shape +# print("Activations shape: ", activations.shape) +# # Augment activations with a column of ones for bias calculation +# ones_col = jnp.ones((B, L, 1)) +# X_tilde = jnp.concatenate([activations, ones_col], axis=-1) # Shape (B, L, D+1) +# d_head = W_Q_a[0].shape[1] +# sqrt_d = jnp.sqrt(float(d_head)) +# C = np.zeros((num_heads, num_heads)) +# for i in range(num_heads): +# # Pre-compute for model A, head i +# tilde_W_Q_a_i = compute_extended_weights(W_Q_a[i], b_Q_a[i]) +# tilde_W_K_a_i = compute_extended_weights(W_K_a[i], b_K_a[i]) +# tilde_W_V_a_i = compute_extended_weights(W_V_a[i], b_V_a[i]) +# Q_a_i = X_tilde @ tilde_W_Q_a_i +# K_a_i = X_tilde @ tilde_W_K_a_i +# S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d +# P_a_i = nn.softmax(S_a_i, axis=-1) +# P_a_i_flat = P_a_i.flatten() +# V_tilde_a_i = X_tilde @ tilde_W_V_a_i +# V_a_i = V_tilde_a_i @ W_O_a[i] +# V_a_i_flat = V_a_i.flatten() +# for j in range(num_heads): +# # Compute for model B, head j +# tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j]) +# tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j]) +# tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j]) +# Q_b_j = X_tilde @ tilde_W_Q_b_j +# K_b_j = X_tilde @ tilde_W_K_b_j +# S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d +# P_b_j = nn.softmax(S_b_j, axis=-1) +# P_b_j_flat = P_b_j.flatten() +# V_tilde_b_j = X_tilde @ tilde_W_V_b_j +# V_b_j = V_tilde_b_j @ W_O_b[j] +# V_b_j_flat = V_b_j.flatten() +# # Cosine similarity for P +# dot_P = jnp.dot(P_a_i_flat, P_b_j_flat) +# norm_P_a = jnp.linalg.norm(P_a_i_flat) +# norm_P_b = jnp.linalg.norm(P_b_j_flat) +# cos_sim_P = dot_P / (norm_P_a * norm_P_b + epsilon) +# cost_P = 1.0 - cos_sim_P +# # Cosine similarity for V +# dot_V = jnp.dot(V_a_i_flat, V_b_j_flat) +# norm_V_a = jnp.linalg.norm(V_a_i_flat) +# norm_V_b = jnp.linalg.norm(V_b_j_flat) +# cos_sim_V = dot_V / (norm_V_a * norm_V_b + epsilon) +# cost_V = 1.0 - cos_sim_V +# total_cost = (alpha * cost_P + (1 - alpha) * cost_V) +# C[i, j] = total_cost +# return C +# Version 3: Normalized Norm +# def compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, +# W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, +# num_heads, activations, alpha=0.5, epsilon=1e-8): +# """ +# Computes the cost matrix for attention head permutation using Normalized Frobenius Norm. +# The cost C[i, j] is the relative error: ||A-B||_F / ||A||_F. +# """ +# # activations: (B, L, D) +# B, L, D = activations.shape +# print("Activations shape: ", activations.shape) +# # Augment activations with a column of ones for bias calculation +# ones_col = jnp.ones((B, L, 1)) +# X_tilde = jnp.concatenate([activations, ones_col], axis=-1) # Shape (B, L, D+1) +# d_head = W_Q_a[0].shape[1] +# sqrt_d = jnp.sqrt(float(d_head)) +# C = np.zeros((num_heads, num_heads)) +# for i in range(num_heads): +# # Pre-compute for model A, head i +# tilde_W_Q_a_i = compute_extended_weights(W_Q_a[i], b_Q_a[i]) +# tilde_W_K_a_i = compute_extended_weights(W_K_a[i], b_K_a[i]) +# tilde_W_V_a_i = compute_extended_weights(W_V_a[i], b_V_a[i]) +# Q_a_i = X_tilde @ tilde_W_Q_a_i +# K_a_i = X_tilde @ tilde_W_K_a_i +# S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d +# P_a_i = nn.softmax(S_a_i, axis=-1) +# V_tilde_a_i = X_tilde @ tilde_W_V_a_i +# V_a_i = V_tilde_a_i @ W_O_a[i] +# for j in range(num_heads): +# # Compute for model B, head j +# tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j]) +# tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j]) +# tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j]) +# Q_b_j = X_tilde @ tilde_W_Q_b_j +# K_b_j = X_tilde @ tilde_W_K_b_j +# S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d +# P_b_j = nn.softmax(S_b_j, axis=-1) +# V_tilde_b_j = X_tilde @ tilde_W_V_b_j +# V_b_j = V_tilde_b_j @ W_O_b[j] +# # Normalized Frobenius norm for P. We flatten the 3D tensor to compute the L2 norm. +# norm_diff_P = jnp.linalg.norm((P_a_i - P_b_j).flatten()) +# norm_P_a = jnp.linalg.norm(P_a_i.flatten()) +# cost_P = norm_diff_P / (norm_P_a + epsilon) +# # Normalized Frobenius norm for V. We flatten the 3D tensor to compute the L2 norm. +# norm_diff_V = jnp.linalg.norm((V_a_i - V_b_j).flatten()) +# norm_V_a = jnp.linalg.norm(V_a_i.flatten()) +# cost_V = norm_diff_V / (norm_V_a + epsilon) +# total_cost = (alpha * cost_P + (1 - alpha) * cost_V) +# C[i, j] = total_cost +# return C +def align_single_head(W_Q_a_i, b_Q_a_i, W_K_a_i, b_K_a_i, W_V_a_i, b_V_a_i, W_O_a_i, + W_Q_b_i, b_Q_b_i, W_K_b_i, b_K_b_i, W_V_b_i, b_V_b_i, W_O_b_i, optimize): + tilde_W_Q_a_i = compute_extended_weights(W_Q_a_i, b_Q_a_i) + tilde_W_K_a_i = compute_extended_weights(W_K_a_i, b_K_a_i) + tilde_W_V_a_i = compute_extended_weights(W_V_a_i, b_V_a_i) + Y_O_a_i = W_O_a_i.T + tilde_W_Q_b_i = compute_extended_weights(W_Q_b_i, b_Q_b_i) + tilde_W_K_b_i = compute_extended_weights(W_K_b_i, b_K_b_i) + tilde_W_V_b_i = compute_extended_weights(W_V_b_i, b_V_b_i) + Y_O_b_i = W_O_b_i.T + A_init = solve_orthogonal(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i) + B_init = solve_orthogonal(Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i) + if optimize: + A, objective_values_A, grad_norms_A, condition_nums_A = optimize_alignment( + A_init, tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i + ) + B, objective_values_B, grad_norms_B, condition_nums_B = optimize_alignment( + B_init, Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i + ) + else: + A = A_init + B = B_init + + A_inv = np.linalg.inv(A) + B_inv = np.linalg.inv(B) + W_Q_aligned = W_Q_b_i @ A.T + b_Q_aligned = b_Q_b_i @ A.T + W_K_aligned = W_K_b_i @ A_inv + b_K_aligned = b_K_b_i @ A_inv + W_V_aligned = W_V_b_i @ B_inv + b_V_aligned = b_V_b_i @ B_inv + W_O_aligned = B @ W_O_b_i + + aligned_params = { + 'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned}, + 'key': {'kernel': W_K_aligned, 'bias': b_K_aligned}, + 'value': {'kernel': W_V_aligned, 'bias': b_V_aligned}, + 'out': {'kernel': W_O_aligned} + } + + if optimize: + return { + 'aligned_params': aligned_params, + 'metrics_A': { + 'objective_values': objective_values_A, + 'grad_norms': grad_norms_A, + 'condition_nums': condition_nums_A + }, + 'metrics_B': { + 'objective_values': objective_values_B, + 'grad_norms': grad_norms_B, + 'condition_nums': condition_nums_B + } + } + return {'aligned_params': aligned_params} + +def apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, h): + aligned_params = {} + for i in range(h): + print(f"Aligning Heads {i}") + tilde_W_Q_a_i = np.vstack([W_Q_a[i], b_Q_a[i].reshape(1, -1)]) + tilde_W_K_a_i = np.vstack([W_K_a[i], b_K_a[i].reshape(1, -1)]) + tilde_W_Q_b_i = np.vstack([W_Q_b[i], b_Q_b[i].reshape(1, -1)]) + tilde_W_K_b_i = np.vstack([W_K_b[i], b_K_b[i].reshape(1, -1)]) + A_i_init = solve_orthogonal(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i) + # A_i = optimize(A_i_init, tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i) + A_i = A_i_init + tilde_W_V_a_i = np.vstack([W_V_a[i], b_V_a[i].reshape(1, -1)]) + tilde_W_V_b_i = np.vstack([W_V_b[i], b_V_b[i].reshape(1, -1)]) + Y_O = W_O_a[i].T + Y_O_prime = W_O_b[i].T + B_i_init = solve_orthogonal(tilde_W_V_a_i, tilde_W_V_b_i, Y_O, Y_O_prime) + # B_i = optimize(B_i_init, tilde_W_V_a_i, tilde_W_V_b_i, Y_O, Y_O_prime) + B_i = B_i_init + A_i_inv = np.linalg.inv(A_i) + B_i_inv = np.linalg.inv(B_i) + + W_Q_aligned = W_Q_b[i] @ A_i.T + b_Q_aligned = b_Q_b[i] @ A_i.T + W_K_aligned = W_K_b[i] @ A_i_inv + b_K_aligned = b_K_b[i] @ A_i_inv + W_V_aligned = W_V_b[i] @ B_i_inv + b_V_aligned = b_V_b[i] @ B_i_inv + W_O_aligned = B_i @ W_O_b[i] + + aligned_params[f'head_{i}'] = { + 'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned}, + 'key': {'kernel': W_K_aligned, 'bias': b_K_aligned}, + 'value': {'kernel': W_V_aligned, 'bias': b_V_aligned}, + 'out': {'kernel': W_O_aligned} + } + return aligned_params + +def merge_aligned_params(aligned_params, h, D, out_bias_b): + query_kernel = np.stack([aligned_params[f'head_{i}']['query']['kernel'] for i in range(h)], axis=1) + query_bias = np.stack([aligned_params[f'head_{i}']['query']['bias'] for i in range(h)], axis=0) + key_kernel = np.stack([aligned_params[f'head_{i}']['key']['kernel'] for i in range(h)], axis=1) + key_bias = np.stack([aligned_params[f'head_{i}']['key']['bias'] for i in range(h)], axis=0) + value_kernel = np.stack([aligned_params[f'head_{i}']['value']['kernel'] for i in range(h)], axis=1) + value_bias = np.stack([aligned_params[f'head_{i}']['value']['bias'] for i in range(h)], axis=0) + out_kernel = np.stack([aligned_params[f'head_{i}']['out']['kernel'] for i in range(h)], axis=0) + # print(query_kernel.shape,"|",query_bias.shape) + # print(key_kernel.shape,"|",key_bias.shape) + # print(value_kernel.shape,"|",value_bias.shape) + # print(out_kernel.shape,) + query_kernel = query_kernel.reshape(-1, D) + key_kernel = key_kernel.reshape(-1, D) + value_kernel = value_kernel.reshape(-1, D) + out_kernel = out_kernel.reshape(D, -1) + query_bias = query_bias.reshape(-1) + key_bias = key_bias.reshape(-1) + value_bias = value_bias.reshape(-1) + return { + 'attention': { + 'query': {'kernel': jnp.array(query_kernel),'bias': jnp.array(query_bias),}, + 'value': {'kernel': jnp.array(value_kernel),'bias': jnp.array(value_bias),}, + 'key': {'kernel': jnp.array(key_kernel),'bias': jnp.array(key_bias),}, + }, + 'output': { + 'dense': {'kernel': jnp.array(out_kernel),'bias': jnp.array(out_bias_b),}, + }, + } + +def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activation, permute_heads=True, optimize=False, alpha=0.5): + attn_a = params_a['vit']['encoder']["layer"][str(layer_idx)]['attention'] + attn_b = params_b['vit']['encoder']["layer"][str(layer_idx)]['attention'] + query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, out_bias_a = extract_attention_params(attn_a) + query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, out_bias_b = extract_attention_params(attn_b) + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_attention_weights(query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, num_heads) + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_attention_weights(query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, num_heads) + if permute_heads: + C = compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, num_heads, activation, alpha) + row_ind, col_ind = linear_sum_assignment(C) + print("Best Permutation Heads:", col_ind) + W_Q_b = [W_Q_b[j] for j in col_ind] + b_Q_b = [b_Q_b[j] for j in col_ind] + W_K_b = [W_K_b[j] for j in col_ind] + b_K_b = [b_K_b[j] for j in col_ind] + W_V_b = [W_V_b[j] for j in col_ind] + b_V_b = [b_V_b[j] for j in col_ind] + W_O_b = [W_O_b[j] for j in col_ind] + if optimize: + metrics_A_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']} + metrics_B_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']} + aligned_params, return_dict = {}, {} + for i in range(num_heads): + result = align_single_head( + W_Q_a[i], b_Q_a[i], W_K_a[i], b_K_a[i], W_V_a[i], b_V_a[i], W_O_a[i], + W_Q_b[i], b_Q_b[i], W_K_b[i], b_K_b[i], W_V_b[i], b_V_b[i], W_O_b[i], optimize + ) + aligned_params[f'head_{i}'] = result['aligned_params'] + if optimize: + for key in metrics_A_all: + metrics_A_all[key].append(result['metrics_A'][key]) + metrics_B_all[key].append(result['metrics_B'][key]) + return_dict['aligned_params'] = merge_aligned_params(aligned_params, num_heads, query_a.shape[1], out_bias_b) + if optimize: + return_dict['metrics_A_all'] = metrics_A_all + return_dict['metrics_B_all'] = metrics_B_all + return return_dict + +def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ + # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), + ] + for name, init_method, permute_heads, optimize in configurations: + aligned_params = copy.deepcopy(params_b) + if optimize: + layer_to_metrics_A = {} + layer_to_metrics_B = {} + for layer_idx in config.lmc_layer_indices: + if activation is not None: + activations_for_layer = activation[layer_idx] + else: activations_for_layer = None + result = align_attention_params( + rng, params_a, aligned_params, layer_idx, config.lmc_config.num_attention_heads, + activations_for_layer, permute_heads=permute_heads, optimize=optimize + ) + aligned_params['vit']['encoder']["layer"][str(layer_idx)]['attention'] = result['aligned_params'] + if optimize: + layer_to_metrics_A[layer_idx] = result['metrics_A_all'] + layer_to_metrics_B[layer_idx] = result['metrics_B_all'] + + total_sum = tree_util.tree_reduce(lambda acc, x: acc + jnp.sum(x), aligned_params, initializer=0) + print(f"{name}: {total_sum}, sanity check") + params_dict[name] = aligned_params + return params_dict + # cost_head = copy.deepcopy(params_b) + # naive_head = copy.deepcopy(params_b) + # for layer_idx in config.lmc_layer_indices: + # aligned_attention_params = cost_align_attention_params(rng, params_a, params_b, str(layer_idx), config.lmc_config.num_attention_heads) + # cost_head['vit']['encoder']["layer"][str(layer_idx)]['attention'] = aligned_attention_params + # for layer_idx in config.lmc_layer_indices: + # aligned_attention_params = naive_align_attention_params(rng, params_a, params_b, str(layer_idx), config.lmc_config.num_attention_heads) + # naive_head['vit']['encoder']["layer"][str(layer_idx)]['attention'] = aligned_attention_params + # return {"cost_head": cost_head, "naive_head": naive_head} diff --git a/src/imagenet/model.py b/src/imagenet/model.py new file mode 100644 index 0000000000000000000000000000000000000000..06bea60f9e60a5888edf436d872f412742c4651a --- /dev/null +++ b/src/imagenet/model.py @@ -0,0 +1,490 @@ +from typing import Optional, Tuple, List +import flax.linen as nn +import jax +import copy +import numpy as np +import jax.numpy as jnp +from flax.linen.attention import dot_product_attention_weights +from flax.core.frozen_dict import FrozenDict, freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ( + FlaxViTPreTrainedModel, + ViTConfig, + FlaxBaseModelOutput, + FlaxBaseModelOutputWithPooling, + FlaxViTPatchEmbeddings, + FlaxViTPooler, + ACT2FN, + FlaxViTIntermediate, + FlaxViTOutput, + FlaxPreTrainedModel, + FlaxSequenceClassifierOutput, +) +from typing import Callable + +def print_model(flax_params, file=None): + flat_params = flatten_dict(flax_params) + for path, value in flat_params.items(): + name = "/".join(path) + line = f"{name} {value.shape}" + if file: + print(line, file=file) + else: + print(line) +def print_model_with_prefix(flax_params, prefix: str, file=None): + flat_params = flatten_dict(flax_params) + for path, value in flat_params.items(): + name = ".".join(path) + if name.startswith(prefix): + line = f"{name} {value.shape} \n {value} \n \n" + if file: print(line, file=file) + else: print(line) + +def create_sinusoidal_positions(n_pos, dim): + position_enc = np.array([[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)]) + sentinel = dim // 2 + dim % 2 + out = np.zeros_like(position_enc) + out[:, 0:sentinel] = np.sin(position_enc[:, 0::2]) + out[:, sentinel:] = np.cos(position_enc[:, 1::2]) + return jnp.array(out) + +class LMCFlaxViTEmbeddings(nn.Module): + """Construct the CLS token, position and patch embeddings.""" + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + self.cls_token = self.param( + "cls_token", + jax.nn.initializers.variance_scaling(self.config.initializer_range**2, "fan_in", "truncated_normal"), + (1, 1, self.config.hidden_size), + ) + self.patch_embeddings = FlaxViTPatchEmbeddings(self.config, dtype=self.dtype) + num_patches = self.patch_embeddings.num_patches + if self.config.position_embeddings == "learnable": + self.position_embeddings = self.param( + "position_embeddings", + jax.nn.initializers.variance_scaling(self.config.initializer_range**2, "fan_in", "truncated_normal"), + (1, num_patches + 1, self.config.hidden_size), + ) + elif self.config.position_embeddings == "sinusoidal": + self.position_embeddings = jnp.expand_dims(create_sinusoidal_positions(num_patches + 1, self.config.hidden_size),axis=0) + self.dropout = nn.Dropout(rate=self.config.hidden_dropout_prob) + def __call__(self, pixel_values, deterministic=True): + batch_size = pixel_values.shape[0] + embeddings = self.patch_embeddings(pixel_values) + cls_tokens = jnp.broadcast_to(self.cls_token, (batch_size, 1, self.config.hidden_size)) + embeddings = jnp.concatenate((cls_tokens, embeddings), axis=1) + if self.config.position_embeddings in ["learnable","sinusoidal"]: + embeddings = embeddings + self.position_embeddings + embeddings = self.dropout(embeddings, deterministic=deterministic) + return embeddings + +class LMCFlaxViTSelfAttention(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + if self.config.hidden_size % self.config.num_attention_heads != 0: + raise ValueError( + "`config.hidden_size`: {self.config.hidden_size} has to be a multiple of `config.num_attention_heads`:" + " {self.config.num_attention_heads}" + ) + self.query = nn.Dense( + self.config.hidden_size, + dtype=self.dtype, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, mode="fan_in", distribution="truncated_normal" + ), + use_bias=self.config.qkv_bias, + ) + self.key = nn.Dense( + self.config.hidden_size, + dtype=self.dtype, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, mode="fan_in", distribution="truncated_normal" + ), + use_bias=self.config.qkv_bias, + ) + self.value = nn.Dense( + self.config.hidden_size, + dtype=self.dtype, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, mode="fan_in", distribution="truncated_normal" + ), + use_bias=self.config.qkv_bias, + ) + if self.config.position_embeddings == "rope": + self.rotary_value = self.config.rotary_value + def __call__(self, hidden_states, sinusoidal_pos, deterministic: bool = True, output_attentions: bool = False): + head_dim = self.config.hidden_size // self.config.num_attention_heads + query_states = self.query(hidden_states).reshape( + hidden_states.shape[:2] + (self.config.num_attention_heads, head_dim) + ) + value_states = self.value(hidden_states).reshape( + hidden_states.shape[:2] + (self.config.num_attention_heads, head_dim) + ) + key_states = self.key(hidden_states).reshape( + hidden_states.shape[:2] + (self.config.num_attention_heads, head_dim) + ) + if sinusoidal_pos is not None and self.config.position_embeddings == 'rope': + if self.rotary_value: + query_states, key_states, value_states = self.apply_rotary_position_embeddings( + sinusoidal_pos, query_states, key_states, value_states + ) + else: + query_states, key_states = self.apply_rotary_position_embeddings( + sinusoidal_pos, query_states, key_states + ) + dropout_rng = None + if not deterministic and self.config.attention_probs_dropout_prob > 0.0: + dropout_rng = self.make_rng("dropout") + attn_weights = dot_product_attention_weights( + query_states, + key_states, + dropout_rng=dropout_rng, + dropout_rate=self.config.attention_probs_dropout_prob, + broadcast_dropout=True, + deterministic=deterministic, + dtype=self.dtype, + precision=None, + ) + attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value_states) + attn_output = attn_output.reshape(attn_output.shape[:2] + (-1,)) + outputs = (attn_output, attn_weights) if output_attentions else (attn_output,) + return outputs + @staticmethod + def apply_rotary_position_embeddings(sinusoidal_pos, query_layer, key_layer, value_layer=None): + sin, cos = jnp.split(sinusoidal_pos, 2, axis=-1) + sin_pos = jnp.stack([sin, sin], axis=-1).reshape(sinusoidal_pos.shape) + cos_pos = jnp.stack([cos, cos], axis=-1).reshape(sinusoidal_pos.shape) + def rotate_layer(layer, sin_pos, cos_pos): + rotate_half_layer = jnp.stack([-layer[..., 1::2], layer[..., ::2]], axis=-1).reshape(layer.shape) + rotary_matrix_cos = jnp.einsum("bslh,...sh->bslh", layer, cos_pos) + rotary_matrix_sin = jnp.einsum("bslh,...sh->bslh", rotate_half_layer, sin_pos) + return rotary_matrix_cos + rotary_matrix_sin + query_layer = rotate_layer(query_layer, sin_pos, cos_pos) + key_layer = rotate_layer(key_layer, sin_pos, cos_pos) + if value_layer is not None: + value_layer = rotate_layer(value_layer, sin_pos, cos_pos) + return query_layer, key_layer, value_layer + return query_layer, key_layer +class LMCFlaxViTSelfOutput(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + self.dense = nn.Dense( + self.config.hidden_size, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, "fan_in", "truncated_normal" + ), + dtype=self.dtype, + ) + self.dropout = nn.Dropout(rate=self.config.hidden_dropout_prob) + + def __call__(self, hidden_states, input_tensor, deterministic: bool = True): + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states, deterministic=deterministic) + return hidden_states +class LMCFlaxViTAttention(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.attention = LMCFlaxViTSelfAttention(self.config, dtype=self.dtype) + self.output = LMCFlaxViTSelfOutput(self.config, dtype=self.dtype) + def __call__(self, hidden_states, sinusoidal_pos, deterministic=True, output_attentions: bool = False): + attn_outputs = self.attention(hidden_states, sinusoidal_pos, deterministic=deterministic, output_attentions=output_attentions) + attn_output = attn_outputs[0] + hidden_states = self.output(attn_output, hidden_states, deterministic=deterministic) + outputs = (hidden_states,) + if output_attentions: + outputs += (attn_outputs[1],) + + return outputs +class LMCFlaxViTMLP(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + self.intermediate = FlaxViTIntermediate(self.config, dtype=self.dtype) + self.output = FlaxViTOutput(self.config, dtype=self.dtype) + def __call__(self, layer_output, attention_output, deterministic: bool = True): + hidden_states = self.intermediate(layer_output) + hidden_states = self.output(hidden_states, attention_output, deterministic=deterministic) + return hidden_states +class LMCFlaxViTLayer(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + self.attention = LMCFlaxViTAttention(self.config, dtype=self.dtype) + self.mlp = LMCFlaxViTMLP(self.config, dtype=self.dtype) + self.layernorm_before = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype) + self.layernorm_after = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype) + def __call__(self, hidden_states,sinusoidal_pos, deterministic: bool = True, output_attentions: bool = False, attention_input:dict|None=None,): + if attention_input is not None: + attention_input = self.layernorm_before(hidden_states) + attention_outputs = self.attention( + attention_input, # self.layernorm_before(hidden_states), # in ViT, layernorm is applied before self-attention + sinusoidal_pos, + deterministic=deterministic, + output_attentions=output_attentions, + ) + attention_output = attention_outputs[0] + # first residual connection + attention_output = attention_output + hidden_states + # in ViT, layernorm is also applied after self-attention + layer_output = self.layernorm_after(attention_output) + hidden_states = self.mlp(layer_output , attention_output, deterministic=deterministic) + outputs = (hidden_states,) + if output_attentions: + outputs += (attention_outputs[1],) + return outputs + +class LMCFlaxViTLayerCollection(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + self.layers = [ + LMCFlaxViTLayer(self.config, name=str(i), dtype=self.dtype) for i in range(self.config.num_hidden_layers) + ] + def __call__( + self, + hidden_states, + sinusoidal_pos, + deterministic: bool = True, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + attention_input:dict|None=None, + ): + all_attentions = () if output_attentions else None + all_hidden_states = () if output_hidden_states else None + + for i, layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + if attention_input is not None: + layer_outputs = layer(hidden_states,sinusoidal_pos,deterministic=deterministic, output_attentions=output_attentions) + else: + attention_input_i = attention_input.get(i, {}) + layer_outputs = layer(hidden_states,sinusoidal_pos,deterministic=deterministic, output_attentions=output_attentions, attention_input=attention_input_i) + attention_input[i]= attention_input_i + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions += (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + outputs = (hidden_states,) + if not return_dict: + return tuple(v for v in outputs if v is not None) + + return FlaxBaseModelOutput( + last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions + ) +class LMCFlaxViTEncoder(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + def setup(self): + num_patches = (self.config.image_size**2)//self.config.patch_size + self.embed_positions = create_sinusoidal_positions( + num_patches + 1, self.config.hidden_size // self.config.num_attention_heads + ) + self.layer = LMCFlaxViTLayerCollection(self.config, dtype=self.dtype) + def __call__( + self, + hidden_states, + deterministic: bool = True, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + attention_input:dict|None=None, + ): + sinusoidal_pos = self.embed_positions[: hidden_states.shape[1], :] + return self.layer( + hidden_states, + sinusoidal_pos, + deterministic=deterministic, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + attention_input=attention_input, + ) + + +class LMCFlaxViTPreTrainedModel(FlaxPreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + config_class = ViTConfig + base_model_prefix = "vit" + main_input_name = "pixel_values" + module_class: nn.Module = None + def __init__( + self, + config: ViTConfig, + input_shape=None, + seed: int = 0, + dtype: jnp.dtype = jnp.float32, + _do_init: bool = True, + **kwargs, + ): + module = self.module_class(config=config, dtype=dtype, **kwargs) + if input_shape is None: + input_shape = (1, config.image_size, config.image_size, config.num_channels) + super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init) + + def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict: + # init input tensors + pixel_values = jnp.zeros(input_shape, dtype=self.dtype) + + params_rng, dropout_rng = jax.random.split(rng) + rngs = {"params": params_rng, "dropout": dropout_rng} + + random_params = self.module.init(rngs, pixel_values, return_dict=False)["params"] + + if params is not None: + random_params = flatten_dict(unfreeze(random_params)) + params = flatten_dict(unfreeze(params)) + for missing_key in self._missing_keys: + params[missing_key] = random_params[missing_key] + self._missing_keys = set() + return freeze(unflatten_dict(params)) + else: + return random_params + + def __call__( + self, + pixel_values, + params: Optional[dict] = None, + dropout_rng: jax.random.PRNGKey = None, + train: bool = False, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + attention_input:dict|None=None, + ): + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + pixel_values = jnp.transpose(pixel_values, (0, 2, 3, 1)) + # Handle any PRNG if needed + rngs = {} + if dropout_rng is not None: + rngs["dropout"] = dropout_rng + + return self.module.apply( + {"params": params or self.params}, + jnp.array(pixel_values, dtype=jnp.float32), + not train, + output_attentions, + output_hidden_states, + return_dict, + rngs=rngs, + attention_input=attention_input, + ) + + +class LMCFlaxViTModule(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 # the dtype of the computation + add_pooling_layer: bool = True + + def setup(self): + self.embeddings = LMCFlaxViTEmbeddings(self.config, dtype=self.dtype) + self.encoder = LMCFlaxViTEncoder(self.config, dtype=self.dtype) + self.layernorm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype) + self.pooler = FlaxViTPooler(self.config, dtype=self.dtype) if self.add_pooling_layer else None + + def __call__( + self, + pixel_values, + deterministic: bool = True, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + attention_input:dict|None=None, + ): + hidden_states = self.embeddings(pixel_values, deterministic=deterministic) + + outputs = self.encoder( + hidden_states, + deterministic=deterministic, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + attention_input=attention_input, + ) + hidden_states = outputs[0] + hidden_states = self.layernorm(hidden_states) + pooled = self.pooler(hidden_states) if self.add_pooling_layer else None + + if not return_dict: + # if pooled is None, don't return it + if pooled is None: + return (hidden_states,) + outputs[1:] + return (hidden_states, pooled) + outputs[1:] + + return FlaxBaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=pooled, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + + +class LMCFlaxViTModel(LMCFlaxViTPreTrainedModel): + module_class = LMCFlaxViTModule + +class LMCFlaxViTForImageClassificationModule(nn.Module): + config: ViTConfig + dtype: jnp.dtype = jnp.float32 + def setup(self): + self.vit = LMCFlaxViTModule(config=self.config, dtype=self.dtype, add_pooling_layer=False) + self.classifier = nn.Dense( + self.config.num_labels, + dtype=self.dtype, + kernel_init=jax.nn.initializers.variance_scaling( + self.config.initializer_range**2, "fan_in", "truncated_normal" + ), + ) + + def __call__( + self, + pixel_values=None, + deterministic: bool = True, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + attention_input:dict|None=None, + ): + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.vit( + pixel_values, + deterministic=deterministic, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + attention_input=attention_input, + ) + hidden_states = outputs[0] + logits = self.classifier(hidden_states[:, 0, :]) + + if not return_dict: + output = (logits,) + outputs[2:] + return output + + return FlaxSequenceClassifierOutput( + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) +class OldLMCFlaxViTForImageClassification(LMCFlaxViTPreTrainedModel): + module_class = LMCFlaxViTForImageClassificationModule + diff --git a/src/imagenet/plot.py b/src/imagenet/plot.py new file mode 100644 index 0000000000000000000000000000000000000000..3c5da19c013281b3875d1f4a5371f1d3cc67f64b --- /dev/null +++ b/src/imagenet/plot.py @@ -0,0 +1,87 @@ +import argparse +import os +import json +import numpy as np +import matplotlib.pyplot as plt + +def calculate_loss_barrier(loss_curve): + return np.max(loss_curve) - (loss_curve[0] + loss_curve[-1]) / 2 + +def main(): + parser = argparse.ArgumentParser(description="Plot ImageNet test loss and accuracy for best weight matching") + parser.add_argument("--file-1", type=str, required=True, help="Path to first JSON results file") + parser.add_argument("--file-2", type=str, required=True, help="Path to second JSON results file") + parser.add_argument("--file-3", type=str, required=True, help="Path to third JSON results file") + parser.add_argument("--output-dir", type=str, default=".", help="Directory to save output plot") + args = parser.parse_args() + + file_paths = [args.file_1, args.file_2, args.file_3] + xtick_labels = [ + ["Model 1", r"$\lambda$", "Model 2"], + ["Model 1", r"$\lambda$", "Model 3"], + ["Model 2", r"$\lambda$", "Model 3"], + ] + + data = [json.load(open(path, 'r')) for path in file_paths] + + num_points = len(data[0]["Naive"]["Test Loss"]) + lambda_values = np.linspace(0, 1, num_points) + + plt.rcParams.update({ + "font.family": "serif", + 'legend.frameon': False, + 'lines.linewidth': 2, + }) + + colors = ["steelblue", "lightsalmon"] + plt.style.use('tableau-colorblind10') + fig, axs = plt.subplots(2, 3, figsize=(13, 6)) + + FONT_SMALL = 11 + FONT_MEDIUM = 13 + FONT_LARGE = 16 + + # Find best index (lowest barrier) for each file + # best_indices = [] + # for d in data: + # best_idx = 0 + # min_barrier = float('inf') + # for idx, loss_curve in enumerate(d["test_loss_interp_clever_list"]): + # barrier = calculate_loss_barrier(loss_curve) + # if barrier < min_barrier: + # min_barrier = barrier + # best_idx = idx + # best_indices.append(best_idx) + + for i in range(3): + d = data[i] + # best_idx = best_indices[i] + + # Row 0: Test Loss + ax = axs[0, i] + ax.plot(lambda_values, d["Naive"]["Test Loss"], label="Naive", color=colors[0]) + ax.plot(lambda_values, d["permu_head_init_ortho_opt"]["Test Loss"], label="WM", color=colors[1]) + ax.set_xticks([0, 0.5, 1]) + ax.set_xticklabels(xtick_labels[i], fontsize=FONT_MEDIUM) + ax.tick_params(axis='y', labelsize=FONT_MEDIUM) + if i == 0: + ax.set_ylabel("Validation Loss", fontsize=FONT_LARGE, labelpad=15) + ax.legend(loc='upper left', fontsize=FONT_SMALL) + + # Row 1: Test Accuracy + ax = axs[1, i] + ax.plot(lambda_values, d["Naive"]["Test Acc"], label="Naive", color=colors[0]) + ax.plot(lambda_values, d["permu_head_init_ortho_opt"]["Test Acc"], label="WM", color=colors[1]) + ax.set_xticks([0, 0.5, 1]) + ax.set_xticklabels(xtick_labels[i], fontsize=FONT_MEDIUM) + ax.tick_params(axis='y', labelsize=FONT_MEDIUM) + if i == 0: + ax.set_ylabel("Validation Accuracy (%)", fontsize=FONT_LARGE, labelpad=15) + + plt.tight_layout(w_pad=4.0, h_pad=2.5) + output_path = os.path.join(args.output_dir, "plot.pdf") if args.output_dir == '.' else args.output_dir + ".pdf" + plt.savefig(output_path) + print(f"Saved plot to {output_path}") + +if __name__ == "__main__": + main() diff --git a/src/imagenet/readme.md b/src/imagenet/readme.md new file mode 100644 index 0000000000000000000000000000000000000000..7f6ced92e125df260ebd8db76098678cb7c556ce --- /dev/null +++ b/src/imagenet/readme.md @@ -0,0 +1,96 @@ +# Linear Mode Connectivity for ViT-MoE on ImageNet + +This document describes the experimental pipeline used to evaluate **Linear Mode Connectivity (LMC)** between independently fine-tuned Vision Transformer (ViT) models augmented with Mixture-of-Experts (MoE) layers on the **ImageNet-1K** dataset. The process includes model training, expert permutation alignment, linear interpolation in parameter space, and visualization of the resulting loss landscape. + +--- + +# Image Classification on ImageNet + +## Download the ImageNet-1K Dataset (ILSVRC 2012) + +To obtain the ImageNet 2012 dataset, navigate to the `/data/` directory and execute the following script: + +```bash +bash src/imagenet/data.sh +``` + +After that, the ImageNet data will be saved in the following directory: + +```bash +/data/imagenet/ +``` + +--- + +## Retrieve Pretrained Vision Transformer (ViT) Weights + +Download the pretrained `ViT-Base-Patch16-224` model by running: + +```bash +CUDA_VISIBLE_DEVICES=0 python src/imagenet/pretrained.py +``` + +The pretrained weights will be saved locally at: + +```bash +./weights/imagenet/vit-base-patch16-224/ +``` + +--- + +## Impact of Feedforward Reinitialization on Pretrained Transformer Performance + +To evaluate the effect of reinitializing the Feedforward Network (FFN) in each Transformer layer, run the following script: + +```bash +CUDA_VISIBLE_DEVICES=0 python src/imagenet/layer_replace_init.py \ + --model-path ./weights/imagenet/vit-base-patch16-224/ \ + --data-path /data/imagenet/ +``` + +This script systematically inserts an MoE block at each layer, fine-tunes the model, and evaluates the resulting performance. + +--- + +## Fine-tune ViT with a Mixture-of-Experts (MoE) Architecture + +To fine-tune a Vision Transformer model with a Mixture-of-Experts (MoE) layer, run: + +```bash +CUDA_VISIBLE_DEVICES=0 python src/imagenet/finetune_moe.py \ + --model-path ./weights/imagenet/vit-base-patch16-224/ \ + --data-path /data/imagenet/ \ + --moe-idx 0 --num-shared-experts 1 \ + --num-routed-experts 7 --topk 2 --seed 0 +``` + +Repeat the command with different seeds (e.g., `--seed 20`, `--seed 40`) to obtain multiple independently fine-tuned models. + +> **Note**: Setting `--topk` equal to `--num-routed-experts` results in a dense MoE configuration. + +--- + +## Expert Matching for Permutation Alignment (Optional but Recommended) + +Due to the permutation invariance of experts in MoE models, aligning expert indices across different runs improves the accuracy and interpretability of LMC analysis. Use the following script to compute a matching between experts from two models: + +```bash +CUDA_VISIBLE_DEVICES=0 python src/imagenet/expert_matching.py \ + --model-a weights/imagenet/finetune/idx0-lr0.0005-seed-0-shared0-routed16-topk2 \ + --model-b weights/imagenet/finetune/idx0-lr0.0005-seed-20-shared0-routed16-topk2 \ + --data-path /data/imagenet/ +``` +--- +## Visualize the Linear Mode Connectivity Curves + +To generate loss curve plots from the interpolation results between matched models, run: + +```bash +python src/imagenet/plot.py \ + --file-1 ./results/imagenet/[idx0-lr0.0005-seed-0-shared0-routed16-topk2+idx0-lr0.0005-seed-20-shared0-routed16-topk2].json \ + --file-2 ./results/imagenet/[idx0-lr0.0005-seed-0-shared0-routed16-topk2+idx0-lr0.0005-seed-40-shared0-routed16-topk2].json \ + --file-3 ./results/imagenet/[idx0-lr0.0005-seed-20-shared0-routed16-topk2+idx0-lr0.0005-seed-40-shared0-routed16-topk2].json \ + --output-dir plots/imagenet/[idx0-lr0.0005-shared0-routed16-topk2] +``` +--- + diff --git a/src/imagenet/test.py b/src/imagenet/test.py new file mode 100644 index 0000000000000000000000000000000000000000..b746441de77088cc1965566b17dadcae38d5c397 --- /dev/null +++ b/src/imagenet/test.py @@ -0,0 +1,305 @@ +import argparse +import os +import jax +import flax +import copy +import torch +import wandb +import optax +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from typing import Any, Dict, List +from flax.jax_utils import replicate, unreplicate +from flax.core.frozen_dict import freeze, unfreeze +from flax.training import checkpoints, train_state +from flax.traverse_util import flatten_dict, unflatten_dict +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from transformers.models.vit.modeling_flax_vit import ViTConfig +from datasets import build_dataset +import multiprocessing as mp +from pprint import pprint +import json +import shutil + +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" + +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +# ---------- Training Utilities ---------- +def get_trainable_mask(params, config): + def is_att_param(keys): + if len(keys) < 5: + return False + return (keys[0] == "vit" and keys[1] == "encoder" and keys[2] == "layer" and int(keys[3]) in config.lmc_layer_indices and keys[4] == "attention") + def label_fn(path, _): + keys = [str(k.key) for k in path] + if is_att_param(keys): return "trainable" + return "frozen" + return jax.tree_util.tree_map_with_path(label_fn, params) +def pretrained2finetune_params(pretrained_params, finetune_params, config): + pretrained_params = unfreeze(pretrained_params) + finetune_params = unfreeze(finetune_params) + # 1. Copy top-level params (embeddings, layernorm, classifier) + finetune_params["vit"]["embeddings"] = copy.deepcopy(pretrained_params["vit"]["embeddings"]) + finetune_params["vit"]["layernorm"] = copy.deepcopy(pretrained_params["vit"]["layernorm"]) + finetune_params["classifier"] = copy.deepcopy(pretrained_params["classifier"]) + # 2. Copy encoder layers + for i in range(config.num_hidden_layers): + str_i = str(i) + if i in config.lmc_layer_indices: + ref_layer = pretrained_params["vit"]["encoder"]["layer"][str_i] # use layer 0 from pretrained + target_layer = finetune_params["vit"]["encoder"]["layer"][str_i] + # Copy shared parts + target_layer["layernorm_before"] = copy.deepcopy(ref_layer["layernorm_before"]) + target_layer["layernorm_after"] = copy.deepcopy(ref_layer["layernorm_after"]) + target_layer["moe"] = copy.deepcopy(ref_layer["moe"]) + # #Copy Attention weights + # target_layer["attention"] = copy.deepcopy(ref_layer["attention"]) + else: + finetune_params["vit"]["encoder"]["layer"][str_i] = copy.deepcopy(pretrained_params["vit"]["encoder"]["layer"][str_i]) + return freeze(finetune_params) + + +# ---------- Main ---------- +def main(args: argparse.Namespace): + if os.path.exists(args.model_path): + config = ViTConfig.from_pretrained(os.path.dirname(args.model_path)) + else: + raise FileNotFoundError(f"Config directory does not exist: {args.model_path}") + os.makedirs(args.wandb_run_dir, exist_ok=True) + # --- Seeds & RNG --- + save_path = os.path.join( + args.save_dir,f"finetune-{config.position_embeddings}-indice{str(args.lmc_layer_indices)}-heads{args.num_attention_heads}" + f"-shared{config.num_shared_experts}-routed{config.num_routed_experts}-topk{config.topk}-seed{args.seed}" + ) + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + pretrained_model = LMCFlaxViTForImageClassification(config, dtype=jnp.dtype(args.dtype)) + lmc_config = copy.deepcopy(config) + lmc_config.num_attention_heads = args.num_attention_heads + lmc_config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_config = lmc_config + config.lmc_layer_indices = args.lmc_layer_indices + pretrained_params = checkpoints.restore_checkpoint(ckpt_dir=args.model_path, target={"params": pretrained_model.params})["params"] + pretrained_model.params = pretrained_params + # --- Initialize fine-tuning model --- + model = LMCFlaxViTForImageClassification(config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype),) + model.config.save_pretrained(save_path) + print_model(model.params) + model.params = pretrained2finetune_params(pretrained_model.params,model.params,config) + model = pretrained_model + label_mask = get_trainable_mask(model.params,config) + print(json.dumps(label_mask, indent=2)) + num_train_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_train_steps, + end_value=args.min_lr, + ) + tx = optax.multi_transform( + transforms={ + 'trainable': optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay + ), + 'frozen': optax.set_to_zero() + }, + param_labels=label_mask + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(ckpt_dir=args.restore_checkpoint_path, target=state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + train_metrics_stack = [] + train_loss = 0.0 + global_step = latest_global_step + best_val_acc1 = 0.0 + ###JUST FOR TESTING#### + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + # val_loss, val_acc1, val_acc5 = float(np.mean(eval_metrics["loss"])), float(np.mean(eval_metrics["acc1"])), float(np.mean(eval_metrics["acc5"])) + print("-" * 100) + print(f"valid loss {val_loss:5.4f} | valid acc@1 {val_acc1:6.2f}% | valid acc@5 {val_acc5:6.2f}%") + print("-" * 100) + exit() + #### START FINETUNING #### + print("Starting Finetuning...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metric, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics_stack.append(train_metric) + train_metrics = get_metrics(train_metrics_stack) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(lambda x: x.mean(), train_metrics) + train_metrics_stack = [] + loss, acc1, acc5 = float(train_metrics["loss"]), float(train_metrics["acc1"]), float(train_metrics["acc5"]) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + # Run eval step + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) # Run train step + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(lambda x: x.mean(), eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 < val_acc1: + best_val_acc1 = val_acc1 + model.params = unreplicate(state).params + best_dir = os.path.join(save_path, f"best_{global_step}") + model.save_pretrained(best_dir) + print(f"Best model saved at step {global_step}") + remove_old_dirs_with_prefix(save_path, "best_", global_step) + # checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1,overwrite=True,orbax_checkpointer=None) + # Save last model + model.params = unreplicate(state).params + last_dir = os.path.join(save_path, f"last_{global_step}") + model.save_pretrained(last_dir) + print(f"Last model saved at step {global_step}") + remove_old_dirs_with_prefix(save_path, "last_", global_step) + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--model-path", type=str, default="", help="Path of Pretrained Model") + parser.add_argument("--num-attention-heads", type=int, default = 1) + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=3, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=8) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default="", help="wandb entity for logging") + parser.add_argument("--wandb-group", default="", help="wandb group for logging") + parser.add_argument("--wandb-project", default="", help="wandb project name for logging") + parser.add_argument("--wandb-run-dir", default=".wandb", help="wandb run dir") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..640c63a771962a9242a118a78eb65d940e0e42ec --- /dev/null +++ b/src/imagenet/train_model.py @@ -0,0 +1,264 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 0 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR10','CIFAR100', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/src/imagenet/transfer.py b/src/imagenet/transfer.py new file mode 100644 index 0000000000000000000000000000000000000000..77464e305fcad537c4b2409ce2f5cfb577507e7f --- /dev/null +++ b/src/imagenet/transfer.py @@ -0,0 +1,186 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from typing import Any, Dict, List +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from flax.training.common_utils import get_metrics, onehot, shard +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from model import OldLMCFlaxViTForImageClassification +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from pprint import pprint +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def pretrained_new2old(new_params,old_params,config): + new_params = unfreeze(new_params) + old_params = unfreeze(old_params) + print_model(new_params) + print_model(old_params) + # 1. Copy top-level params (embeddings, layernorm, classifier) + new_params["vit"]["embeddings"] = copy.deepcopy(old_params["vit"]["embeddings"]) + new_params["vit"]["layernorm"] = copy.deepcopy(old_params["vit"]["layernorm"]) + new_params["classifier"] = copy.deepcopy(old_params["classifier"]) + # 2. Copy encoder layers + for i in range(config.num_hidden_layers): + str_i = str(i) + ref_layer = old_params["vit"]["encoder"]["layer"][str_i] # use layer 0 from pretrained + target_layer = new_params["vit"]["encoder"]["layer"][str_i] + # Copy shared parts + target_layer["layernorm_before"] = copy.deepcopy(ref_layer["layernorm_before"]) + target_layer["layernorm_after"] = copy.deepcopy(ref_layer["layernorm_after"]) + target_layer["moe"]['shared_experts']['intermediate'] = copy.deepcopy(ref_layer["mlp"]['intermediate']['dense']) + target_layer["moe"]['shared_experts']['output'] = copy.deepcopy(ref_layer["mlp"]['output']['dense']) + #Copy Attention weights + target_layer["attention"] = copy.deepcopy(ref_layer["attention"]) + return freeze(new_params) +def main(args: argparse.Namespace): + train_loader, val_loader = imagenet_data_loader(args) + save_path = "/mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/" + old_config = ViTConfig.from_pretrained('/mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/config.json') + old_model = OldLMCFlaxViTForImageClassification(old_config,dtype=jnp.bfloat16) + new_config = copy.deepcopy(old_config) + new_config.position_embeddings = old_config.position_embeddings + new_config.rotary_value = old_config.rotary_value + new_config.q_lora_rank = 8 + new_config.qk_rope_head_dim = 64 + new_config.kv_lora_rank = 8 + new_config.v_head_dim = 64 + new_config.qk_nope_head_dim = 64 + new_config.attention_bias = True + new_config.routed_scaling_factor = 1.0 + new_config.lmc_layer_indices = [] + new_config.num_shared_experts = 1 + new_config.num_routed_experts = 0 + new_config.topk = 0 + # 1. Recreate LR schedule + lr_schedule = optax.warmup_cosine_decay_schedule( + init_value=1e-6, + peak_value=5e-4, + warmup_steps=5*5004, + decay_steps=300*5004, + end_value=1e-5, + ) + # 2. Recreate optimizer tx + tx = optax.adamw( + learning_rate=lr_schedule, + b1=0.9, + b2=0.999, + eps=1e-8, + weight_decay=0.01, + ) + old_state = train_state.TrainState.create(apply_fn=old_model.__call__, params=old_model.params, tx=tx) + old_params = old_state.params + new_model = LMCFlaxViTForImageClassification(new_config,input_shape=(1,new_config.image_size, new_config.image_size, new_config.num_channels),seed=args.seed,dtype=jnp.bfloat16) + new_model.params = pretrained_new2old(new_params=copy.deepcopy(new_model.params),old_params=copy.deepcopy(old_state.params),config=new_model.config) + # print(new_model.params['classifier']['bias']) + # 3. Create new TrainState with new_model + new_state = train_state.TrainState.create(apply_fn=new_model.__call__,params=new_model.params,tx=tx,) + print(old_state.params['classifier']['bias']) + print(new_state.params['classifier']['bias']) + # 4. Load old values + os.makedirs(save_path,exist_ok=True) + new_model.config.save_pretrained(save_path) + # checkpoints.save_checkpoint(ckpt_dir=save_path,target=new_state,step=old_state.step,prefix="best_",keep=1,overwrite=True) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + state = replicate(new_state) + rng = jax.random.PRNGKey(0) + train_metrics_stack = [] + train_loss = 0.0 + best_val_acc1 = 0.0 + ###JUST FOR TESTING#### + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + # Run eval step + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) # Run train step + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(lambda x: x.mean(), eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + print("-" * 100) + print(f"valid loss {val_loss:5.4f} | valid acc@1 {val_acc1:6.2f}% | valid acc@5 {val_acc5:6.2f}%") + print("-" * 100) +if __name__ =="__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + parser.add_argument("--batch-size", type=int, default=256) + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=8) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/src/imagenet/transform.py b/src/imagenet/transform.py new file mode 100644 index 0000000000000000000000000000000000000000..892a8ddcdaf979003b15f18e5ec6e860aebc1b37 --- /dev/null +++ b/src/imagenet/transform.py @@ -0,0 +1,83 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +def main(args: argparse.Namespace): + save_path = "/mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0" + model = LMCFlaxViTForImageClassification.from_pretrained("/mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/temp_1291032") + step = 1291032 + num_global_steps = 5004*args.epochs + num_warmup_steps = 5004*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + state = state.replace(step=step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=state,step=state.step,prefix="last_",keep=1) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=3, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/src/imagenet/utils.py b/src/imagenet/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0f425c2e9fad683b43d981a2458ed6aa5973dbcd --- /dev/null +++ b/src/imagenet/utils.py @@ -0,0 +1,118 @@ +import operator +import re +import time +from contextlib import contextmanager + +import jax.numpy as jnp +from flax import traverse_util +from flax.core import freeze, unfreeze +from jax import random +from jax.tree_util import tree_reduce, tree_map + +rngmix = lambda rng, x: random.fold_in(rng, hash(x) % 2**16) + +@contextmanager +def timeblock(name): + start = time.time() + try: + yield + finally: + end = time.time() + print(f"{name} took {end - start:.5f} seconds") + +class RngPooper: + """A stateful wrapper around stateless random.PRNGKey's.""" + + def __init__(self, init_rng): + self.rng = init_rng + + def poop(self): + self.rng, rng_key = random.split(self.rng) + return rng_key + +def l1prox(x, alpha): + return jnp.sign(x) * jnp.maximum(0, jnp.abs(x) - alpha) + +def ec2_get_instance_type(): + # See also https://stackoverflow.com/questions/51486405/aws-ec2-command-line-display-instance-type/51486782 + return open("/sys/devices/virtual/dmi/id/product_name").read().strip() + +# Utilities for dealing with flax model parameters +def partition(pred, iterable): + trues = [] + falses = [] + for item in iterable: + if pred(item): + trues.append(item) + else: + falses.append(item) + return trues, falses + +def partition_dict(pred, d): + trues = {} + falses = {} + for k, v in d.items(): + if pred(k): + trues[k] = v + else: + falses[k] = v + return trues, falses + +def flatten_params(params): + return {"/".join(k): v for k, v in traverse_util.flatten_dict(unfreeze(params)).items()} + +def unflatten_params(flat_params): + return freeze( + traverse_util.unflatten_dict({tuple(k.split("/")): v + for k, v in flat_params.items()})) + +def merge_params(a, b): + return unflatten_params({**a, **b}) + +def kmatch(pattern, key): + regex = "^" + i = 0 + while i < len(pattern): + if pattern[i] == "*": + if i + 1 < len(pattern) and pattern[i + 1] == "*": + regex += "(.*)" + i += 2 + else: + regex += "([^\/]*)" + i += 1 + else: + regex += pattern[i] + i += 1 + regex += "$" + return re.fullmatch(regex, key) + +assert kmatch("*", "a") is not None +assert kmatch("*", "a").group(0) == "a" +assert kmatch("*", "a").group(1) == "a" +assert kmatch("abc", "def") is None +assert kmatch("abc/*/ghi", "abc/def/ghi").group(1) == "def" +assert kmatch("abc/**/jkl", "abc/def/ghi/jkl").group(1) == "def/ghi" +assert kmatch("abc/*/jkl", "abc/def/ghi/jkl") is None +assert kmatch("**/*", "abc/def/ghi/jkl").group(1) == "abc/def/ghi" +assert kmatch("**/*", "abc/def/ghi/jkl").group(2) == "jkl" + +def lerp(lam, t1, t2): + return tree_map(lambda a, b: (1 - lam) * a + lam * b, t1, t2) + +def tree_norm(t): + return jnp.sqrt(tree_reduce(operator.add, tree_map(lambda x: jnp.sum(x**2), t))) + +def tree_l2(t1, t2): + return tree_norm(tree_map(lambda x, y: x - y, t1, t2)) + +def slerp(lam, t1, t2): + # See https://en.wikipedia.org/wiki/Slerp + om = jnp.arccos( + tree_reduce(operator.add, tree_map(lambda x, y: jnp.sum(x * y), t1, t2)) / + (tree_norm(t1) * tree_norm(t2))) + sinom = jnp.sin(om) + return tree_map( + lambda x, y: jnp.sin((1 - lam) * om) / sinom * x + jnp.sin(lam * om) / sinom * y, + t1, + t2, + ) diff --git a/src/lgmodeling/__pycache__/data_utils.cpython-310.pyc b/src/lgmodeling/__pycache__/data_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33984f38485a32254fbf866ffb6488a36ffdec4f Binary files /dev/null and b/src/lgmodeling/__pycache__/data_utils.cpython-310.pyc differ diff --git a/src/lgmodeling/__pycache__/lmc_model.cpython-310.pyc b/src/lgmodeling/__pycache__/lmc_model.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8fa101907aa7af109b9c2fa1a21171ef6f608834 Binary files /dev/null and b/src/lgmodeling/__pycache__/lmc_model.cpython-310.pyc differ diff --git a/src/lgmodeling/__pycache__/matching_utils.cpython-310.pyc b/src/lgmodeling/__pycache__/matching_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..087a2630f51afab1e9cdca35f765493200979fb8 Binary files /dev/null and b/src/lgmodeling/__pycache__/matching_utils.cpython-310.pyc differ diff --git a/src/lgmodeling/__pycache__/model.cpython-310.pyc b/src/lgmodeling/__pycache__/model.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d8ed2c324f5aedc2af07c3efbbb33eedbdd400cc Binary files /dev/null and b/src/lgmodeling/__pycache__/model.cpython-310.pyc differ diff --git a/src/lgmodeling/__pycache__/utils.cpython-310.pyc b/src/lgmodeling/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a103647707e7c2c65d2e5d7c29cce2d03442407 Binary files /dev/null and b/src/lgmodeling/__pycache__/utils.cpython-310.pyc differ diff --git a/src/lgmodeling/__pycache__/vocabulary.cpython-310.pyc b/src/lgmodeling/__pycache__/vocabulary.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..261baf2ea147b25061dc83f3fb89054831a3db00 Binary files /dev/null and b/src/lgmodeling/__pycache__/vocabulary.cpython-310.pyc differ diff --git a/src/lgmodeling/att_matching.py b/src/lgmodeling/att_matching.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f30c46e662f760d73249ae84b3d5648a358351 --- /dev/null +++ b/src/lgmodeling/att_matching.py @@ -0,0 +1,258 @@ +import argparse +import os +import jax +import copy +import json +import torch +import optax +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from jax import random, vmap +from typing import Any, Dict, List +from flax.core import freeze, unfreeze +from flax.serialization import from_bytes +from utils import flatten_params, lerp, unflatten_params +from flax.traverse_util import flatten_dict, unflatten_dict +from lmc_model import LMCFlaxGPT2LMHeadModel, print_model +from transformers.models.gpt2.modeling_flax_gpt2 import GPT2Config, FlaxGPT2LMHeadModel +from flax.training import checkpoints, train_state +from flax.training.common_utils import get_metrics, onehot, shard +from matching_utils import weight_matching_attn, all_matching_attn +from data_utils import get_lm_corpus +import matplotlib.pyplot as plt +import numpy as np + +# def check_params_nan(params, name="params"): +# flat_params = flatten_dict(unfreeze(params)) +# found_nan = False +# for path, val in flat_params.items(): +# if isinstance(val, jnp.ndarray): +# if jnp.isnan(val).any(): +# print(f"🚫 NaN detected in {name}: {'/'.join(path)} | shape={val.shape}") +# found_nan = True +# if not found_nan: +# print(f"✅ No NaNs found in {name}.") +# else: +# print(f"⚠️ NaNs found in {name}.") + + +def load_flax_params(checkpoint_dir, target): + """ + Loads Flax parameters from a msgpack file. + + Args: + checkpoint_dir (str): path to folder containing flax_model.msgpack + target (PyTree): a target object matching the structure of your model's parameters + + Returns: + PyTree: deserialized parameters + """ + msgpack_path = os.path.join(checkpoint_dir, "flax_model.msgpack") + + with open(msgpack_path, "rb") as f: + packed_bytes = f.read() + + params = from_bytes(target, packed_bytes) + return params +def prepare_lm_batch(data: torch.Tensor, target: torch.Tensor) -> Dict[str, Any]: + """ + Convert and shard a language modeling batch from PyTorch to JAX. + Args: + data (torch.Tensor): Input data of shape (seq_len, batch) + target (torch.Tensor): Target data of shape (seq_len, batch) + Returns: + Dict[str, jnp.ndarray]: Dict with 'data' and 'target', both sharded + with shape (n_devices, batch_per_device, seq_len) + """ + # Transpose to (batch, seq_len), then convert to jnp arrays + input_ids = jnp.array(data.T) + target = jnp.array(target.T) + # Shard across devices + return {'input_ids': input_ids,'target': target} +def make_stuff(model): + apply_fn = model.__call__ + @jax.jit + def batch_eval(params, batch): + labels = batch.pop("target") # destructive, just like your train_step + logits = apply_fn(**batch, params=params, train=False)[0] + loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])).mean() + return logits, labels, loss + @jax.jit + def step(state: train_state.TrainState, batch, dropout_rng): + dropout_rng, new_dropout_rng = jax.random.split(dropout_rng, 2) + def loss_fn(params): + labels = batch.pop("target") + logits = apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0] + loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + new_state = state.apply_gradients(grads=grads) + metrics = {"batch_loss": loss,"logits": logits,} + return new_state, metrics, new_dropout_rng + def dataset_loss_and_ppl(params, dataloader): + """ + Iterate once over `dataloader`, pplumulate token-level CE, and + return (mean_loss, perplexity). Works on **one device**. + """ + total_loss, total_tok = 0.0, 0 + pbar = tqdm(dataloader, desc="Evaluating", leave=False) + for eval_data, eval_target, _ in pbar: + eval_batch = prepare_lm_batch(eval_data, eval_target) + _, labels, loss = batch_eval(params,eval_batch) + ntok = jnp.sum(labels != -100) + total_loss += loss * ntok + total_tok += ntok + pbar.set_postfix(loss=f"{loss:.4f}", ppl=f"{np.exp(loss):.2f}") + mean_loss = (total_loss / total_tok).item() + ppl = jnp.exp(mean_loss).item() + return mean_loss, ppl + return {"batch_eval": batch_eval,"step": step,"dataset_loss_and_ppl": dataset_loss_and_ppl,} +def compute_interpolation(params_a, params_b_target, lambdas, stuff, val_ds, test_ds, desc="Interpolation"): + train_loss_interp, test_loss_interp = [], [] + train_ppl_interp, test_ppl_interp = [], [] + for lam in tqdm(lambdas, desc=desc): + p_interp = freeze(lerp(lam, unfreeze(params_a), unfreeze(params_b_target))) + train_loss, train_ppl = stuff["dataset_loss_and_ppl"](p_interp, val_ds) + # test_loss, test_ppl = stuff["dataset_loss_and_ppl"](p_interp, test_ds) + test_loss, test_ppl = train_loss, train_ppl + train_loss_interp.append(train_loss) + test_loss_interp.append(test_loss) + train_ppl_interp.append(train_ppl) + test_ppl_interp.append(test_ppl) + return { + "Val Loss": [float(f"{x:.4f}") for x in train_loss_interp], + "Test Loss": [float(f"{x:.4f}") for x in test_loss_interp], + "Val PPL": [float(f"{x:.4f}") for x in train_ppl_interp], + "Test PPL": [float(f"{x:.4f}") for x in test_ppl_interp] + } +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-a", type=str, required=True, help="Path to first fine-tuned GPT2 model checkpoint") + parser.add_argument("--model-b", type=str, required=True, help="Path to second fine-tuned GPT@ model checkpoint") + parser.add_argument("--seed", type=int, default =0) + parser.add_argument("--data-path", type=str, default="./data/lm1b", help="train datset paths (multiple paths)") + parser.add_argument('--dataset', type=str, default='lm1b',choices=['wt103', 'lm1b', 'enwik8', 'text8'],help='dataset name') + parser.add_argument("--batch-size", type=int, default=24, help="train, eval batch size (batch size will be devided by device count)") + parser.add_argument('--tgt_len', type=int, default=256,help='number of tokens to predict') + parser.add_argument('--eval_tgt_len', type=int, default=256,help='number of tokens to predict for evaluation') + parser.add_argument('--ext_len', type=int, default=0,help='length of the extended context') + parser.add_argument('--mem_len', type=int, default=0,help='length of the retained previous heads') + parser.add_argument("--save-path", type=str, default="/", help="Path to plot directory") + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + args = parser.parse_args() + corpus = get_lm_corpus(args.data_path, args.dataset) + ntokens = len(corpus.vocab) + eval_batch_size = 12 + tr_iter = corpus.get_iterator('train', args.batch_size, args.tgt_len, ext_len=args.ext_len) + va_iter = corpus.get_iterator('valid', eval_batch_size, args.eval_tgt_len, ext_len=args.ext_len) + te_iter = corpus.get_iterator('test', eval_batch_size, args.eval_tgt_len, ext_len=args.ext_len) + val_ds = va_iter + test_ds = te_iter + config = GPT2Config.from_json_file(os.path.join(os.path.dirname(args.model_a).rstrip("/"),'config.json')) + lmc_config = GPT2Config(**config.lmc_config) + config.lmc_config = lmc_config + model = LMCFlaxGPT2LMHeadModel(config,input_shape=(1, args.tgt_len),seed=args.seed,dtype=jnp.dtype(args.dtype),) + print_model(model.params) + if os.path.exists(args.model_a) and os.path.exists(args.model_b): + params_a = load_flax_params(args.model_a,copy.deepcopy(model.params)) + # check_params_nan(params_a, name="params_a") + params_b = load_flax_params(args.model_b,copy.deepcopy(model.params)) + # check_params_nan(params_b, name="params_b") + else: + raise FileNotFoundError(f"Checkpoint path does not exist") + stuff = make_stuff(model = model) + lambdas = jnp.linspace(0, 1, num=25) + rng = random.PRNGKey(args.seed) + # Compute naive interpolation + naive_results = compute_interpolation(params_a, params_b, lambdas, stuff, val_ds, test_ds, desc="Naive Interpolation") + print(json.dumps({"Naive": naive_results}, indent=2)) + all_results = {"Naive": naive_results} + # Compute weight matching interpolations for each method + aligned_models = weight_matching_attn(rng, params_a, params_b, None, config) + for method, params_b_aligned in aligned_models.items(): + method_results = compute_interpolation(params_a, params_b_aligned, lambdas, stuff, val_ds, test_ds, desc=f"{method} Interpolation") + all_results[method] = method_results + print(json.dumps({method: method_results}, indent=2)) + # permutations_results = {} + # aligned_models = all_matching_attn(rng, params_a, params_b, config) + # for method, params_b_aligned in aligned_models.items(): + # method_results = compute_interpolation(params_a, params_b_aligned, lambdas, stuff, val_ds, test_ds, desc=f"{method} Interpolation") + # permutations_results[method] = method_results + # print(json.dumps({method: method_results}, indent=2)) + # Save directories + os.makedirs(f"./plots/{args.dataset}", exist_ok=True) + os.makedirs(f"./results/{args.dataset}", exist_ok=True) + # Save results JSON + print("Save List of Values...") + name_a = os.path.basename(os.path.dirname(args.model_a).rstrip("/")) + name_b = os.path.basename(os.path.dirname(args.model_b).rstrip("/")) + result_path = f'results/{args.dataset}/[{name_a}+{name_b}].json' + with open(result_path, 'w') as f: + json.dump(all_results, f, indent=2) + # permute_path = f'results/{args.dataset}/permute[{name_a}+{name_b}].json' + # with open(permute_path, 'w') as f: + # json.dump(permutations_results, f, indent=2) + # Plot + # print("Generating plots...") + # plot_path = f"./plots/{args.dataset}/[{name_a}+{name_b}].pdf" + # plt.rcParams.update({ + # "font.family": "serif", + # 'legend.frameon': False, + # 'lines.linewidth': 2, + # 'font.size': 13, + # 'axes.labelsize': 16, + # 'xtick.labelsize': 11, + # 'ytick.labelsize': 11, + # 'legend.fontsize': 11, + # }) + # plt.style.use('tableau-colorblind10') + # num_points = len(all_results["Naive"]["Val Loss"]) + # lambda_values = np.linspace(0, 1, num_points) + # fig, axs = plt.subplots(1, 2, figsize=(12, 5)) # 1 row, 2 columns + # metrics = ["Val Loss", "Test Loss"] + # custom_colors = [ + # "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", + # "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf" + # ] + # for idx, metric in enumerate(metrics): + # ax = axs[idx] + # for j, method in enumerate(all_results): # Naive + aligned methods + # ax.plot(lambda_values, + # all_results[method][metric], + # label=method, + # color=custom_colors[j % len(custom_colors)]) + # ax.set_xticks([0, 0.5, 1]) + # ax.set_xticklabels(["Model 1", r"$\lambda$", "Model 2"]) + # ax.set_ylabel(metric) + # ax.legend(loc='best') + + # plt.tight_layout() + # plt.savefig(plot_path.replace(".pdf", "_row0.pdf")) + # plt.close() + # fig, axs = plt.subplots(2, 2, figsize=(12, 10)) + # metrics = ["Val Loss", "Test Loss", "Val Loss", "Test Loss"] + # positions = [(0, 0), (0, 1), (1, 0), (1, 1)] + # custom_colors = [ + # "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", + # "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf" + # ] + # for metric, pos in zip(metrics, positions): + # row, col = pos + # ax = axs[row, col] + # if row == 0: + # for idx, method in enumerate(all_results): + # ax.plot(lambda_values, all_results[method][metric], label=method,color=custom_colors[idx % len(custom_colors)]) + # if row == 1: + # for idx, method in enumerate(permutations_results): + # ax.plot(lambda_values, permutations_results[method][metric], label=method) + # ax.set_xticks([0, 0.5, 1]) + # ax.set_xticklabels(["Model 1", r"$\lambda$", "Model 2"]) + # ax.set_ylabel(metric) + # ax.legend(loc='best') + # plt.tight_layout() + # plt.savefig(plot_path) + # plt.close() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/lgmodeling/data_utils.py b/src/lgmodeling/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d16cb0e5f2b909b2fe9304c6e96fbd55ce45a1c9 --- /dev/null +++ b/src/lgmodeling/data_utils.py @@ -0,0 +1,343 @@ +import os +import glob + +import numpy as np +import torch + +from vocabulary import Vocab + + + +class LMOrderedIterator(object): + def __init__(self, data, bsz, bptt, device='cpu', ext_len=None): + """ + data -- LongTensor -- the LongTensor is strictly ordered + """ + self.bsz = bsz + self.bptt = bptt + self.ext_len = ext_len if ext_len is not None else 0 + + self.device = device + + # Work out how cleanly we can divide the dataset into bsz parts. + self.n_step = data.size(0) // bsz + + # Trim off any extra elements that wouldn't cleanly fit (remainders). + data = data.narrow(0, 0, self.n_step * bsz) + + # Evenly divide the data across the bsz batches. + self.data = data.view(bsz, -1).t().contiguous().to(device) + + # Number of mini-batches + self.n_batch = (self.n_step + self.bptt - 1) // self.bptt + + def get_batch(self, i, bptt=None): + if bptt is None: + bptt = self.bptt + seq_len = min(bptt, self.data.size(0)-1-i) + + end_idx = i + seq_len + beg_idx = max(0, i-self.ext_len) + + data = self.data[beg_idx:end_idx] + target = self.data[i+1:i+1+seq_len] + + return data, target, seq_len + + def get_fixlen_iter(self, start=0): + for i in range(start, self.data.size(0) - 1, self.bptt): + yield self.get_batch(i) + + def get_varlen_iter(self, start=0, std=5, min_len=5, max_deviation=3): + max_len = self.bptt + max_deviation * std + i = start + while True: + bptt = self.bptt if np.random.random() < 0.95 else self.bptt / 2. + bptt = min(max_len, max(min_len, int(np.random.normal(bptt, std)))) + data, target, seq_len = self.get_batch(i, bptt) + i += seq_len + yield data, target, seq_len + if i >= self.data.size(0) - 2: + break + + def __iter__(self): + return self.get_fixlen_iter() + + +class LMShuffledIterator(object): + def __init__(self, data, bsz, bptt, device='cpu', ext_len=None, + shuffle=False): + """ + data -- list[LongTensor] -- there is no order among the LongTensors + """ + self.data = data + + self.bsz = bsz + self.bptt = bptt + self.ext_len = ext_len if ext_len is not None else 0 + + self.device = device + self.shuffle = shuffle + + def get_sent_stream(self): + # index iterator + if self.shuffle: + epoch_indices = np.random.permutation(len(self.data)) + else: + epoch_indices = np.array(range(len(self.data))) + + # sentence iterator + for idx in epoch_indices: + yield self.data[idx] + + def stream_iterator(self, sent_stream): + # streams for each data in the batch + streams = [None] * self.bsz + + data = torch.LongTensor(self.bptt, self.bsz) + target = torch.LongTensor(self.bptt, self.bsz) + + n_retain = 0 + + while True: + # data : [n_retain+bptt x bsz] + # target : [bptt x bsz] + data[n_retain:].fill_(-1) + target.fill_(-1) + + valid_batch = True + + for i in range(self.bsz): + n_filled = 0 + try: + while n_filled < self.bptt: + if streams[i] is None or len(streams[i]) <= 1: + streams[i] = next(sent_stream) + # number of new tokens to fill in + n_new = min(len(streams[i]) - 1, self.bptt - n_filled) + # first n_retain tokens are retained from last batch + data[n_retain+n_filled:n_retain+n_filled+n_new, i] = \ + streams[i][:n_new] + target[n_filled:n_filled+n_new, i] = \ + streams[i][1:n_new+1] + streams[i] = streams[i][n_new:] + n_filled += n_new + except StopIteration: + valid_batch = False + break + + if not valid_batch: + return + + data = data.to(self.device) + target = target.to(self.device) + + yield data, target, self.bptt + + n_retain = min(data.size(0), self.ext_len) + if n_retain > 0: + data[:n_retain] = data[-n_retain:] + data.resize_(n_retain + self.bptt, data.size(1)) + + def __iter__(self): + # sent_stream is an iterator + sent_stream = self.get_sent_stream() + + for batch in self.stream_iterator(sent_stream): + yield batch + + +class LMMultiFileIterator(LMShuffledIterator): + def __init__(self, paths, vocab, bsz, bptt, device='cpu', ext_len=None, + shuffle=False): + + self.paths = paths + self.vocab = vocab + + self.bsz = bsz + self.bptt = bptt + self.ext_len = ext_len if ext_len is not None else 0 + + self.device = device + self.shuffle = shuffle + + def get_sent_stream(self, path): + sents = self.vocab.encode_file(path, add_double_eos=True) + if self.shuffle: + np.random.shuffle(sents) + sent_stream = iter(sents) + + return sent_stream + + def __iter__(self): + if self.shuffle: + np.random.shuffle(self.paths) + + for path in self.paths: + # sent_stream is an iterator + sent_stream = self.get_sent_stream(path) + for batch in self.stream_iterator(sent_stream): + yield batch + + +class LMSlidingWindowIterator(object): + def __init__(self, data, bsz, bptt, device='cpu', ext_len=None): + """ + data -- LongTensor -- the LongTensor is strictly ordered + """ + self.bsz = bsz + self.bptt = bptt + self.ext_len = ext_len if ext_len is not None else 0 + + self.device = device + + # Work out how cleanly we can divide the dataset into bsz parts. + self.n_step = data.size(0) // bsz + self.has_remainder = (data.size(0) % bsz != 0) + if self.has_remainder: + print(f"LMSlidingWindowIterator, Warning: trimming eval text! " + f"data size: {data.size(0)} , batch size: {bsz}") + + # Trim off any extra elements that wouldn't cleanly fit (remainders). + main_data = data.narrow(dim=0, start=0, length=self.n_step * bsz) + # Evenly divide the data across the bsz batches. + self.data = main_data.view(bsz, -1).t().contiguous().to(device) + + # Number of mini-batches + self.n_batch = (self.n_step + self.bptt - 1) // self.bptt + + def get_batch(self, i, bptt=None): + if bptt is None: + bptt = self.bptt + + seq_len = min(bptt, self.data.size(0)-1-i) + assert seq_len == bptt + + end_idx = i + seq_len + beg_idx = max(0, i-self.ext_len) + + data = self.data[beg_idx:end_idx] + target = self.data[i+1:i+1+seq_len] + + return data, target, seq_len + + def get_fixlen_iter(self, start=0): + # first position of the segment is not predicted, thus -1. + for i in range(start, self.data.size(0) - self.bptt): + yield self.get_batch(i) + + def __iter__(self): + return self.get_fixlen_iter() + + def __len__(self): + return self.data.size(0) - self.bptt + 1 + + +class Corpus(object): + def __init__(self, path, dataset, *args, **kwargs): + self.dataset = dataset + self.vocab = Vocab(*args, **kwargs) + + if self.dataset in ['ptb', 'wt2', 'enwik8', 'text8']: + self.vocab.count_file(os.path.join(path, 'train.txt')) + self.vocab.count_file(os.path.join(path, 'valid.txt')) + self.vocab.count_file(os.path.join(path, 'test.txt')) + + elif self.dataset == 'wt103': + self.vocab.count_file(os.path.join(path, 'train.txt')) + elif self.dataset == 'lm1b': + train_path_pattern = os.path.join( + path, '1-billion-word-language-modeling-benchmark-r13output', + 'training-monolingual.tokenized.shuffled', 'news.en-*') + train_paths = glob.glob(train_path_pattern) + # the vocab will load from file when build_vocab() is called + + self.vocab.build_vocab() + + if self.dataset in ['ptb', 'wt2', 'wt103']: + self.train = self.vocab.encode_file( + os.path.join(path, 'train.txt'), ordered=True) + self.valid = self.vocab.encode_file( + os.path.join(path, 'valid.txt'), ordered=True) + self.test = self.vocab.encode_file( + os.path.join(path, 'test.txt'), ordered=True) + + elif self.dataset in ['enwik8', 'text8']: + self.train = self.vocab.encode_file( + os.path.join(path, 'train.txt'), ordered=True, add_eos=False) + self.valid = self.vocab.encode_file( + os.path.join(path, 'valid.txt'), ordered=True, add_eos=False) + self.test = self.vocab.encode_file( + os.path.join(path, 'test.txt'), ordered=True, add_eos=False) + + elif self.dataset == 'lm1b': + self.train = train_paths + self.valid = self.vocab.encode_file( + os.path.join(path, 'valid.txt'), ordered=False, + add_double_eos=True) + self.test = self.vocab.encode_file( + os.path.join(path, 'test.txt'), ordered=False, + add_double_eos=True) + + def get_iterator(self, split, *args, sliding_window=False, **kwargs): + if split == 'train': + if self.dataset in ['ptb', 'wt2', 'wt103', 'enwik8', 'text8']: + data_iter = LMOrderedIterator(self.train, *args, **kwargs) + elif self.dataset == 'lm1b': + kwargs['shuffle'] = True + data_iter = LMMultiFileIterator( + self.train, self.vocab, *args, **kwargs) + + elif split in ['valid', 'test']: + data = self.valid if split == 'valid' else self.test + if self.dataset in ['ptb', 'wt2', 'wt103', 'enwik8', 'text8']: + if sliding_window: + data_iter = LMSlidingWindowIterator(data, *args, **kwargs) + else: + data_iter = LMOrderedIterator(data, *args, **kwargs) + elif self.dataset == 'lm1b': + data_iter = LMShuffledIterator(data, *args, **kwargs) + + return data_iter + + +def get_lm_corpus(datadir, dataset): + fn = os.path.join(datadir, 'cache.pt') + if os.path.exists(fn): + print('Loading cached dataset...') + corpus = torch.load(fn, weights_only=False) + else: + print('Producing dataset {}...'.format(dataset)) + kwargs = {} + if dataset in ['wt103', 'wt2']: + kwargs['special'] = [''] + kwargs['lower_case'] = False + elif dataset == 'ptb': + kwargs['special'] = [''] + kwargs['lower_case'] = True + elif dataset == 'lm1b': + kwargs['special'] = [] + kwargs['lower_case'] = False + kwargs['vocab_file'] = os.path.join(datadir, '1b_word_vocab.txt') + elif dataset in ['enwik8', 'text8']: + pass + + corpus = Corpus(datadir, dataset, **kwargs) + torch.save(corpus, fn) + + return corpus + + +if __name__ == '__main__': + import argparse + parser = argparse.ArgumentParser(description='unit test') + parser.add_argument('--datadir', type=str, default='../data/text8', + help='location of the data corpus') + parser.add_argument('--dataset', type=str, default='text8', + choices=['ptb', 'wt2', 'wt103', 'lm1b', 'enwik8', + 'text8'], + help='dataset name') + args = parser.parse_args() + + corpus = get_lm_corpus(args.datadir, args.dataset) + print('Vocab size : {}'.format(len(corpus.vocab.idx2sym))) \ No newline at end of file diff --git a/src/lgmodeling/finetune.py b/src/lgmodeling/finetune.py new file mode 100644 index 0000000000000000000000000000000000000000..71c0b0c95a4211dae1ada698a3b8d2000ef28fb6 --- /dev/null +++ b/src/lgmodeling/finetune.py @@ -0,0 +1,359 @@ +import os +import math +import time +import copy +import json +import torch +import jax +import argparse +import optax +import wandb +import shutil +import itertools +from tqdm import tqdm +from typing import Any, Dict, List +from copy import deepcopy +from datasets import Dataset +from datetime import timedelta +from data_utils import get_lm_corpus +from flax.jax_utils import replicate, unreplicate +from flax.core.frozen_dict import freeze, unfreeze +from flax.training import train_state, checkpoints +from flax.traverse_util import flatten_dict, unflatten_dict +from flax.training.common_utils import get_metrics, onehot, shard +from transformers.models.gpt2.modeling_flax_gpt2 import GPT2Config +from lmc_model import LMCFlaxGPT2LMHeadModel, print_model +from data_utils import get_lm_corpus +import jax.numpy as jnp +os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" +os.environ["NVIDIA_TF32_OVERRIDE"] = "0" +os.environ["JAX_DEFAULT_MATMUL_PRECISION"] = "highest" +os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) + + + +def prepare_lm_batch(data: torch.Tensor, target: torch.Tensor) -> Dict[str, Any]: + """ + Convert and shard a language modeling batch from PyTorch to JAX. + Args: + data (torch.Tensor): Input data of shape (seq_len, batch) + target (torch.Tensor): Target data of shape (seq_len, batch) + Returns: + Dict[str, jnp.ndarray]: Dict with 'data' and 'target', both sharded + with shape (n_devices, batch_per_device, seq_len) + """ + # Transpose to (batch, seq_len), then convert to jnp arrays + input_ids = jnp.array(data.T) + target = jnp.array(target.T) + # Shard across devices + return {'input_ids': shard(input_ids),'target': shard(target)} +def decay_mask_fn(params): + flat_params = flatten_dict(params) + flat_mask = { + path: (path[-1] != "bias" and path[-2:] not in [("ln_1", "scale"), ("ln_2", "scale"), ("ln_f", "scale")]) + for path in flat_params + } + return unflatten_dict(flat_mask) +# ---------- Training Utilities ---------- +def get_trainable_mask(params, config): + def is_trainable_param(keys): + # Match MoE parameters in transformer/h/{moe_idx}/mlp/(gate|routed_experts_*) + if len(keys) < 4: return False + if keys[0] == "transformer" and keys[1] == "h" and int(keys[2]) in config.lmc_layer_indices and keys[3] == "attn": return True + if config.finetune_mlp == True and keys[0] == "transformer" and keys[1] == "h" and int(keys[2]) in config.lmc_layer_indices and keys[3] == "moe": return True + return False + def label_fn(path, _): + keys = [str(k.key) for k in path] + return "trainable" if is_trainable_param(keys) else "frozen" + return jax.tree_util.tree_map_with_path(label_fn, params) +def pretrained2finetune_params(pretrained_params, finetune_params, config): + pretrained_params = unfreeze(pretrained_params) + finetune_params = unfreeze(finetune_params) + # 1. Copy top-level embeddings and final layer norm + finetune_params["transformer"]["wte"] = copy.deepcopy(pretrained_params["transformer"]["wte"]) + if config.position_embeddings == "learnable": + finetune_params["transformer"]["wpe"] = copy.deepcopy(pretrained_params["transformer"]["wpe"]) + finetune_params["transformer"]["ln_f"] = copy.deepcopy(pretrained_params["transformer"]["ln_f"]) + # 2. Copy encoder layers + for i in range(config.num_hidden_layers): + str_i = str(i) + if i in config.lmc_layer_indices: + # Handle MoE layer: copy attention and norms from pretrained + ref_layer = copy.deepcopy(pretrained_params["transformer"]["h"][str_i]) + target_layer = copy.deepcopy(finetune_params["transformer"]["h"][str_i]) + target_layer["ln_1"] = copy.deepcopy(ref_layer["ln_1"]) + # target_layer["attn"] = copy.deepcopy(ref_layer["attn"]) + target_layer["ln_2"] = copy.deepcopy(ref_layer["ln_2"]) + if config.finetune_mlp == False: + target_layer["moe"] = copy.deepcopy(ref_layer["moe"]) + else: + # Standard block, copy all directly + finetune_params["transformer"]["h"][str_i] = copy.deepcopy(pretrained_params["transformer"]["h"][str_i]) + return freeze(finetune_params) + +def main(args: argparse.Namespace): + if os.path.exists(args.model_path): + config = GPT2Config.from_pretrained(os.path.dirname(args.model_path)) + pretrained_model = LMCFlaxGPT2LMHeadModel(config) + else: + raise FileNotFoundError(f"Config directory does not exist: {os.path.dirname(args.model_path)}") + os.makedirs(args.wandb_run_dir, exist_ok=True) + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + name=f"finetune-{config.position_embeddings}-indice{','.join(str(i) for i in args.lmc_layer_indices)}-heads{args.n_head}" + f"-shared{config.num_shared_experts}-routed{config.num_routed_experts}-topk{config.topk}-mlp{str(args.finetune_mlp)}-seed{args.seed}", + save_code=True + ) + save_path = os.path.join(args.model_save_dir,wandb.run.name) + wandb.config = dict(vars(args)) + corpus = get_lm_corpus(args.data_path, args.dataset) + ntokens = len(corpus.vocab) + args.n_token = ntokens + eval_batch_size = 12 + tr_iter = corpus.get_iterator('train', args.batch_size, args.tgt_len, ext_len=args.ext_len) + va_iter = corpus.get_iterator('valid', eval_batch_size, args.eval_tgt_len, ext_len=args.ext_len) + te_iter = corpus.get_iterator('test', eval_batch_size, args.eval_tgt_len, ext_len=args.ext_len) + lmc_config = copy.deepcopy(config) + lmc_config.n_head = args.n_head + lmc_config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_config = lmc_config + config.lmc_layer_indices = args.lmc_layer_indices + config.finetune_mlp = args.finetune_mlp + pretrained_params = checkpoints.restore_checkpoint(ckpt_dir=args.model_path, target={"params": pretrained_model.params})["params"] + pretrained_model.params = pretrained_params + # --- Initialize fine-tuning model --- + model = LMCFlaxGPT2LMHeadModel(config,input_shape=(1, args.tgt_len),seed=args.seed,dtype=jnp.dtype(args.dtype),) + model.config.save_pretrained(save_path) + print_model(model.params) + model.params = pretrained2finetune_params(pretrained_model.params,model.params,config) + # model = pretrained_model + label_mask = get_trainable_mask(model.params,config) + print(json.dumps(label_mask, indent=2)) + num_train_steps = args.max_step + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=0.0, + peak_value=args.learning_rate, + warmup_steps=args.warmup_step, + decay_steps=args.max_step, + end_value=args.eta_min, + ) + tx = optax.multi_transform( + transforms={ + 'trainable': optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay_rate + ), + 'frozen': optax.set_to_zero() + }, + param_labels=label_mask + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + latest_train_step = state.step + def train_step(state, batch, dropout_rng): + dropout_rng, new_dropout_rng = jax.random.split(dropout_rng) + def loss_fn(params): + labels = batch.pop("target") + logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0] + loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])).mean() + return loss + grad_fn = jax.value_and_grad(loss_fn) + loss, grads = grad_fn(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + new_state = state.apply_gradients(grads=grads) + metrics = {"loss": loss,"learning_rate": lr_schedule(state.step)} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return new_state, metrics, new_dropout_rng + def eval_step(state, batch): + labels = batch.pop("target") + logits = model(**batch, params=state.params, train=False)[0] + loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])).mean() + metrics = {"eval_loss": loss} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + state = replicate(state) + rng = jax.random.PRNGKey(args.seed) + train_metrics_stack = [] + train_step = int(jax.device_get(unreplicate(state.step))) + train_loss = 0.0 + best_val_loss = float("inf") + log_start_time = time.time() + eval_start_time = time.time() + # ###JUST FOR TESTING#### + eval_results = [] + for eval_data, eval_target, _ in va_iter: + eval_batch = prepare_lm_batch(eval_data, eval_target) + eval_metric = parallel_eval_step(state, eval_batch) + eval_results.append(eval_metric) + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(lambda x: x.mean(), eval_metrics) + val_loss = float(eval_metrics["eval_loss"]) + val_ppl = math.exp(val_loss) + print("-" * 100) + print( + f"| Eval {train_step // args.eval_frequency:3d} at step {train_step:8d} | " + f"time: {time.time() - eval_start_time:5.2f}s | " + f"valid loss {val_loss:5.2f} | valid ppl {val_ppl:9.3f}" + ) + print("-" * 100) + # #### START FINETUNING #### + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in itertools.count(start=1): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_iter = tr_iter.get_varlen_iter() if getattr(args, "varlen", False) else tr_iter + train_metrics_stack = [] + for batch_idx, (data, target, seq_len) in enumerate(tqdm(train_iter)): + if train_step >= args.max_step: + break + # Prepare and shard batch + batch = prepare_lm_batch(data, target) + # Run train step + state, train_metric, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics_stack.append(train_metric) + train_step += 1 + # Logging + if train_step % args.logging_frequency == 0: + train_metrics = get_metrics(train_metrics_stack) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(lambda x: x.mean(), train_metrics) + train_metrics_stack = [] + loss = float(train_metrics["loss"]) + ppl = math.exp(loss) + bpc = loss/math.log(2) + curr_lr = float(lr_schedule(train_step)) + elapsed = time.time() - log_start_time + if(args.dataset in ["wt103","lm1b"]): + print( + f"| epoch {epoch:3d} step {train_step:8d} | " + f"{batch_idx+1:6d} batches | lr {curr_lr:.3g} " + f"| ms/batch {elapsed * 1000 / args.logging_frequency:5.2f} | " + f"loss {loss:5.2f} | ppl {ppl:9.3f}" + ) + wandb.log({"loss": loss,"ppl": ppl,"learning_rate": curr_lr}, step=train_step) + elif(args.dataset in ["enwik8","text8"]): + print( + f"| epoch {epoch:3d} step {train_step:8d} | " + f"{batch_idx+1:6d} batches | lr {curr_lr:.3g} " + f"| ms/batch {elapsed * 1000 / args.logging_frequency:5.2f} | " + f"loss {loss:5.2f} | bpc {bpc:9.3f}" + ) + wandb.log({"loss": loss,"bpc": bpc,"learning_rate": curr_lr}, step=train_step) + log_start_time = time.time() + # Evaluation + if train_step % args.eval_frequency == 0: + eval_results = [] + for eval_data, eval_target, _ in va_iter: + eval_batch = prepare_lm_batch(eval_data, eval_target) + eval_metric = parallel_eval_step(state, eval_batch) + eval_results.append(eval_metric) + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(lambda x: x.mean(), eval_metrics) + + val_loss = float(eval_metrics["eval_loss"]) + val_ppl = math.exp(val_loss) + val_bpc = val_loss/math.log(2) + print("-" * 100) + if(args.dataset in ["wt103","lm1b"]): + print( + f"| Eval {train_step // args.eval_frequency:3d} at step {train_step:8d} | " + f"time: {time.time() - eval_start_time:5.2f}s | " + f"valid loss {val_loss:5.2f} | valid ppl {val_ppl:9.3f}" + ) + wandb.log({"eval_loss": val_loss,"eval_ppl": val_ppl}, step=train_step) + elif(args.dataset in ["enwik8","text8"]): + print( + f"| Eval {train_step // args.eval_frequency:3d} at step {train_step:8d} | " + f"time: {time.time() - eval_start_time:5.2f}s | " + f"valid loss {val_loss:5.2f} | valid bpc {val_bpc:9.3f}" + ) + wandb.log({"eval_loss": val_loss,"eval_bpc": val_bpc}, step=train_step) + print("-" * 100) + # Save best checkpoint + if val_loss < best_val_loss: + best_val_loss = val_loss + model.params = unreplicate(state).params + best_dir = os.path.join(save_path, f"best_{train_step}") + model.save_pretrained(best_dir) + print(f"✅ Best model saved at step {train_step}") + remove_old_dirs_with_prefix(save_path, "best_", train_step) + eval_start_time = time.time() + # Periodic checkpoint + if train_step % args.save_frequency == 0: + model.params = unreplicate(state).params + last_dir = os.path.join(save_path, f"last_{train_step}") + model.save_pretrained(last_dir) + print(f"💾 Checkpoint saved at step {train_step}") + remove_old_dirs_with_prefix(save_path, "last_", train_step) + # checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=int(jax.device_get(unreplicate(state.step))),prefix="last_",keep=1,overwrite=True) + print(f"Checkpoint saved at step {save_path}") + if train_step >= args.max_step: + print("-" * 100) + print("End of training") + break + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--model-path", type=str, default="", help="Path of Pretrained Model") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--rotary-dim", type=int,default=0,help=' Number of dimensions in the embedding that Rotary Position Embedding is applied to..') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--n_layer', type=int, default=12,help='number of total layers') + parser.add_argument('--n_head', type=int, default=10,help='number of heads') + parser.add_argument('--n_embd', type=int, default=500,help='model dimension') + parser.add_argument('--n_inner', type=int, default=1000,help='inner dimension in FF') + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument("--finetune-mlp",action="store_true",help="Enable fine-tuning for the MLP. Default is False.") + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--seed", type=int, default=0, help="random seed for RNG state") + parser.add_argument("--data-path", type=str, default="/cm/archive/vinhbk1/datasets/wikitext103", help="train datset paths (multiple paths)") + parser.add_argument('--dataset', type=str, default='wt103',choices=['wt103', 'lm1b', 'enwik8', 'text8'],help='dataset name') + parser.add_argument('--max_step', type=int, default=500000,help='upper epoch limit') + parser.add_argument('--warmup_step', type=int, default=2000,help='upper epoch limit') + parser.add_argument("--batch-size", type=int, default=96, help="train, eval batch size (batch size will be devided by device count)") + parser.add_argument('--tgt_len', type=int, default=256,help='number of tokens to predict') + parser.add_argument('--eval_tgt_len', type=int, default=256,help='number of tokens to predict for evaluation') + parser.add_argument('--ext_len', type=int, default=0,help='length of the extended context') + parser.add_argument('--mem_len', type=int, default=0,help='length of the retained previous heads') + parser.add_argument("--learning-rate", type=float, default=0.00025, help="learning rate") + parser.add_argument("--weight-decay-rate", type=float, default=0.01, help="weight deacy rate for lr scheduler") + parser.add_argument('--eta_min', type=float, default=1.0e-8,help='min learning rate for cosine scheduler') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default="", help="wandb entity for logging") + parser.add_argument("--wandb-group", default="", help="wandb group for logging") + parser.add_argument("--wandb-project", default="GPT2-Wikitext103", help="wandb project name for logging") + parser.add_argument("--wandb-run-dir", default=".wandb", help="wandb run dir") + parser.add_argument("--logging-frequency", type=int, default=200, help="do logging every logging_frequency step") + parser.add_argument("--eval-frequency", type=int, default=4000, help="do evalution every eval_frequency step") + parser.add_argument("--save-frequency", type=int, default=4000, help="do saving checkpoint every save_frequencey step") + parser.add_argument("--model-save-dir", type=str, default="artifacts/", help="checkpoint saving dir") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + main(parser.parse_args()) \ No newline at end of file diff --git a/src/lgmodeling/getdata.py b/src/lgmodeling/getdata.py new file mode 100644 index 0000000000000000000000000000000000000000..257b15631e8f47f732c004414555ce0d182ded83 --- /dev/null +++ b/src/lgmodeling/getdata.py @@ -0,0 +1,23 @@ +from datasets import load_dataset +dataset = load_dataset("wikitext", "wikitext-103-v1") + +train_file = '/mnt/data/vinhbk/wikitext103/train.txt' +valid_file = '/mnt/data/vinhbk/wikitext103/valid.txt' +test_file = '/mnt/data/vinhbk/wikitext103/test.txt' + +# Save the train split +with open(train_file, 'w', encoding='utf-8') as f: + for item in dataset['train']: + f.write(item['text'] + '\n') + +# Save the validation split +with open(valid_file, 'w', encoding='utf-8') as f: + for item in dataset['validation']: + f.write(item['text'] + '\n') + +# Save the test split +with open(test_file, 'w', encoding='utf-8') as f: + for item in dataset['test']: + f.write(item['text'] + '\n') + +print("Files saved successfully!") \ No newline at end of file diff --git a/src/lgmodeling/lmc_model.py b/src/lgmodeling/lmc_model.py new file mode 100644 index 0000000000000000000000000000000000000000..02e6237886cb0cff42949d7b53ae691fcbdd4c21 --- /dev/null +++ b/src/lgmodeling/lmc_model.py @@ -0,0 +1,783 @@ +from typing import Optional, Tuple, List +import flax.linen as nn +import jax +import numpy as np +import jax.numpy as jnp +from jax import lax +from flax.core.frozen_dict import FrozenDict, freeze, unfreeze +from flax.linen.attention import dot_product_attention_weights +from flax.traverse_util import flatten_dict, unflatten_dict +from flax.linen import combine_masks, make_causal_mask +from transformers.models.gpt2.modeling_flax_gpt2 import ( + FlaxConv1D, + FlaxGPT2MLP, + GPT2Config, + FlaxBaseModelOutputWithPastAndCrossAttentions, + FlaxCausalLMOutputWithCrossAttentions, +) +from transformers.modeling_flax_utils import ACT2FN, FlaxPreTrainedModel +from typing import Callable, Literal +def print_model(flax_params, file=None): + flat_params = flatten_dict(flax_params) + for path, value in flat_params.items(): + name = "/".join(path) + if hasattr(value, "shape"): + line = f"{name} {value.shape}" + else: + line = f"{name} {type(value)}" + if file: + print(line, file=file) + else: + print(line) + +def create_sinusoidal_positions(num_pos, dim): + inv_freq = 1.0 / (10000 ** (np.arange(0, dim, 2) / dim)) + sinusoid_inp = np.einsum("i , j -> i j", np.arange(num_pos), inv_freq).astype("float32") + sin, cos = np.sin(sinusoid_inp), np.cos(sinusoid_inp) + + sentinel = dim // 2 + dim % 2 + out = np.zeros((num_pos, dim)) + out[:, 0:sentinel] = sin + out[:, sentinel:] = cos + + return jnp.array(out) + + +def rotate_every_two(tensor): + rotate_half_tensor = jnp.stack((-tensor[:, :, :, 1::2], tensor[:, :, :, ::2]), axis=-1) + rotate_half_tensor = rotate_half_tensor.reshape(rotate_half_tensor.shape[:-2] + (-1,)) + return rotate_half_tensor + + +def apply_rotary_pos_emb(tensor, sincos): + sin_pos, cos_pos = sincos + sin_pos = sin_pos[:, :, None, :].repeat(2, 3) + cos_pos = cos_pos[:, :, None, :].repeat(2, 3) + return (tensor * cos_pos) + (rotate_every_two(tensor) * sin_pos) + + +class LMCFlaxGPT2Attention(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + causal: bool = True + is_cross_attention: bool = False + def setup(self): + config = self.config + self.embed_dim = config.hidden_size + self.num_heads = config.n_head + self.head_dim = self.embed_dim // self.num_heads + if self.config.position_embeddings == "rope": + self.rotary_dim = self.head_dim + if self.is_cross_attention: + self.c_attn = FlaxConv1D(2 * self.embed_dim, dtype=self.dtype) + self.q_attn = FlaxConv1D(self.embed_dim, dtype=self.dtype) + else: + self.c_attn = FlaxConv1D(3 * self.embed_dim, dtype=self.dtype) + self.c_proj = FlaxConv1D(self.embed_dim, dtype=self.dtype) + self.resid_dropout = nn.Dropout(rate=config.resid_pdrop) + if self.causal: + self.causal_mask = make_causal_mask( + jnp.ones((1, config.max_position_embeddings), dtype="bool"), dtype="bool" + ) + if self.config.position_embeddings == "rope": + pos_embd_dim = self.rotary_dim or self.embed_dim + self.embed_positions = create_sinusoidal_positions(config.max_position_embeddings, pos_embd_dim) + def _split_heads(self, hidden_states): + return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads, self.head_dim)) + def _merge_heads(self, hidden_states): + return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,)) + @nn.compact + def _concatenate_to_cache(self, key, value, query, attention_mask): + """ + This function takes projected key, value states from a single input token and concatenates the states to cached + states from previous steps. This function is slightly adapted from the official Flax repository: + https://github.com/google/flax/blob/491ce18759622506588784b4fca0e4bf05f8c8cd/flax/linen/attention.py#L252 + """ + # detect if we're initializing by absence of existing cache data. + is_initialized = self.has_variable("cache", "cached_key") + cached_key = self.variable("cache", "cached_key", jnp.zeros, key.shape, key.dtype) + cached_value = self.variable("cache", "cached_value", jnp.zeros, value.shape, value.dtype) + cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32)) + if is_initialized: + *batch_dims, max_length, num_heads, depth_per_head = cached_key.value.shape + # update key, value caches with our new 1d spatial slices + cur_index = cache_index.value + indices = (0,) * len(batch_dims) + (cur_index, 0, 0) + key = lax.dynamic_update_slice(cached_key.value, key, indices) + value = lax.dynamic_update_slice(cached_value.value, value, indices) + cached_key.value = key + cached_value.value = value + num_updated_cache_vectors = query.shape[1] + cache_index.value = cache_index.value + num_updated_cache_vectors + # causal mask for cached decoder self-attention: our single query position should only attend to those key positions that have already been generated and cached, not the remaining zero elements. + pad_mask = jnp.broadcast_to( + jnp.arange(max_length) < cur_index + num_updated_cache_vectors, + tuple(batch_dims) + (1, num_updated_cache_vectors, max_length), + ) + attention_mask = combine_masks(pad_mask, attention_mask) + return key, value, attention_mask + + def __call__( + self, + hidden_states, + key_value_states: Optional[jnp.ndarray] = None, + attention_mask=None, + position_ids=None, + deterministic: bool = True, + init_cache: bool = False, + output_attentions: bool = False, + ): + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + batch_size = hidden_states.shape[0] + if not is_cross_attention: + qkv_out = self.c_attn(hidden_states) + query, key, value = jnp.split(qkv_out, 3, axis=2) + else: + q_out = self.q_attn(hidden_states) + (query,) = jnp.split(q_out, 1, axis=2) + kv_out = self.c_attn(key_value_states) + key, value = jnp.split(kv_out, 2, axis=2) + query = self._split_heads(query) + key = self._split_heads(key) + value = self._split_heads(value) + if self.config.position_embeddings == "rope": + sincos = jnp.take(self.embed_positions, position_ids, axis=0) + sincos = jnp.split(sincos, 2, axis=-1) + if self.rotary_dim is not None: + k_rot = key[:, :, :, : self.rotary_dim] + k_pass = key[:, :, :, self.rotary_dim :] + + q_rot = query[:, :, :, : self.rotary_dim] + q_pass = query[:, :, :, self.rotary_dim :] + + k_rot = apply_rotary_pos_emb(k_rot, sincos) + q_rot = apply_rotary_pos_emb(q_rot, sincos) + + key = jnp.concatenate([k_rot, k_pass], axis=-1) + query = jnp.concatenate([q_rot, q_pass], axis=-1) + else: + key = apply_rotary_pos_emb(key, sincos) + query = apply_rotary_pos_emb(query, sincos) + query_length, key_length = query.shape[1], key.shape[1] + if self.causal: + if self.has_variable("cache", "cached_key"): + mask_shift = self.variables["cache"]["cache_index"] + max_decoder_length = self.variables["cache"]["cached_key"].shape[1] + causal_mask = lax.dynamic_slice( + self.causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length) + ) + else: + causal_mask = self.causal_mask[:, :, :query_length, :key_length] + causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:]) + # combine masks if needed + if attention_mask is not None and self.causal: + attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape) + attention_mask = combine_masks(attention_mask, causal_mask) + elif self.causal: + attention_mask = causal_mask + elif attention_mask is not None: + attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2)) + dropout_rng = None + if not deterministic and self.config.attn_pdrop > 0.0: + dropout_rng = self.make_rng("dropout") + # During fast autoregressive decoding, we feed one position at a time, + # and cache the keys and values step by step. + if self.causal and (self.has_variable("cache", "cached_key") or init_cache): + key, value, attention_mask = self._concatenate_to_cache(key, value, query, attention_mask) + # transform boolean mask into float mask + if attention_mask is not None: + attention_bias = lax.select( + attention_mask > 0, + jnp.full(attention_mask.shape, 0.0).astype(self.dtype), + jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype), + ) + else: + attention_bias = None + # usual dot product attention + attn_weights = dot_product_attention_weights( + query, key, + bias=attention_bias, + dropout_rng=dropout_rng, + dropout_rate=self.config.attn_pdrop, + deterministic=deterministic, + dtype=self.dtype, + precision=None, + ) + attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value) + attn_output = self._merge_heads(attn_output) + attn_output = self.c_proj(attn_output) + attn_output = self.resid_dropout(attn_output, deterministic=deterministic) + outputs = (attn_output, attn_weights) if output_attentions else (attn_output,) + return outputs +class LMCFlaxGPT2Router(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.n_routed_experts = self.config.num_routed_experts + self.n_group = getattr(self.config, "n_group", 1) + self.topk_group = getattr(self.config, "topk_group", 1) + self.top_k = self.config.topk + self.routed_scaling_factor = self.config.routed_scaling_factor + self.norm_topk_prob = getattr(self.config, "norm_topk_prob", False) + # Weight and bias for router computation + kernel_init = jax.nn.initializers.normal(self.config.initializer_range) + self.router_weight = self.param( + "router_weight", kernel_init, (self.n_routed_experts, self.config.hidden_size) + ) + self.router_bias = self.param( + "router_bias", lambda rng, shape: jnp.zeros(shape, dtype=self.dtype), (self.n_routed_experts,) + ) + self.e_score_correction_bias = self.param( + "e_score_correction_bias", lambda rng, shape: jnp.zeros(shape, dtype=self.dtype), (self.n_routed_experts,) + ) + def get_topk_indices(self, scores): + scores_for_choice = scores + self.e_score_correction_bias[None, :] + scores_grouped = scores_for_choice.reshape( + -1, self.n_group, self.n_routed_experts // self.n_group + ) + top2_scores = jax.lax.top_k(scores_grouped, 2)[0] + group_scores = jnp.sum(top2_scores, axis=-1) + top_group_scores, group_idx = jax.lax.top_k(group_scores, self.topk_group) + group_mask = jnp.zeros_like(group_scores) + group_mask = group_mask.at[ + jnp.arange(group_mask.shape[0])[:, None], group_idx + ].set(1) + group_mask_expanded = jnp.repeat( + group_mask[:, :, None], + self.n_routed_experts // self.n_group, + axis=-1, + ) + score_mask = group_mask_expanded.reshape( + -1, self.n_routed_experts + ) + + scores_for_choice = jnp.where( + score_mask, + scores_for_choice, + jnp.zeros_like(scores_for_choice), + ) + + topk_weights, topk_indices = jax.lax.top_k(scores_for_choice, self.top_k) + + return topk_indices, topk_weights + + def __call__(self, hidden_states): + router_logits = jnp.matmul(hidden_states, self.router_weight.T) + self.router_bias + scores = jax.nn.sigmoid(router_logits) + + if self.n_group > 1: + topk_indices, topk_weights = self.get_topk_indices(scores) + else: + topk_weights, topk_indices = jax.lax.top_k(scores, self.top_k) + + if self.norm_topk_prob: + denominator = jnp.sum(topk_weights, axis=-1, keepdims=True) + 1e-20 + topk_weights = topk_weights / denominator + + topk_weights = topk_weights * self.routed_scaling_factor + + return topk_indices, topk_weights +class LMCFlaxGPT2MoE(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + def setup(self): + self.n_experts = self.config.num_routed_experts + self.n_shared_experts = self.config.num_shared_experts + self.top_k = self.config.topk + self.router = LMCFlaxGPT2Router(config=self.config, dtype=self.dtype) + self.experts = [ + FlaxGPT2MLP( + config=self.config, + intermediate_size=self.config.n_inner, + dtype=self.dtype, + ) + for _ in range(self.n_experts) + ] + self.shared_experts = FlaxGPT2MLP( + config=self.config, + intermediate_size=self.config.n_inner* self.n_shared_experts, + dtype=self.dtype, + ) + + def __call__(self, hidden_states, deterministic: bool = True): + residual = hidden_states + orig_shape = hidden_states.shape + if self.n_experts == 0 or self.top_k == 0: + hidden_states = self.shared_experts(hidden_states, deterministic=deterministic) + return hidden_states + hidden_states_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + topk_indices, topk_weights = self.router(hidden_states_flat) + expert_outputs = [] + for expert in self.experts: + output = expert(hidden_states_flat, deterministic=deterministic) # [n_tokens, hidden_dim] + expert_outputs.append(output) + expert_outputs = jnp.stack(expert_outputs, axis=1) # [n_tokens, n_experts, hidden_dim] + # Build routing mask: + routing_mask = jax.nn.one_hot(topk_indices, self.n_experts, dtype=self.dtype) + routing_mask = routing_mask.sum(axis=1) > 0 # [n_tokens, n_experts] + # Compute weights per expert + weights_per_expert = jax.nn.one_hot(topk_indices, self.n_experts, dtype=self.dtype) + weights_per_expert = (weights_per_expert * topk_weights[..., None]).sum(axis=1) + # Mask out experts not routed + weights_per_expert = weights_per_expert * routing_mask.astype(self.dtype) + # Multiply and sum + weighted_expert_outputs = expert_outputs * weights_per_expert[..., None] # [n_tokens, n_experts, hidden_dim] + final_output = weighted_expert_outputs.sum(axis=1) # [n_tokens, hidden_dim] + # Add shared expert + shared_output = self.shared_experts(residual, deterministic=deterministic) + final_output = final_output.reshape(orig_shape) + shared_output + return final_output +class LMCFlaxGPT2Block(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + def setup(self): + hidden_size = self.config.hidden_size + inner_dim = self.config.n_inner if self.config.n_inner is not None else 4 * hidden_size + self.ln_1 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype) + self.attn = LMCFlaxGPT2Attention(self.config, dtype=self.dtype) + self.ln_2 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype) + if self.config.add_cross_attention: + self.crossattention = LMCFlaxGPT2Attention( + config=self.config, dtype=self.dtype, causal=False, is_cross_attention=True + ) + self.ln_cross_attn = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype) + self.moe = LMCFlaxGPT2MoE(self.config, dtype=self.dtype) + def __call__( + self, + hidden_states, + attention_mask=None, + position_ids=None, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + deterministic: bool = True, + init_cache: bool = False, + output_attentions: bool = False, + ): + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + attn_outputs = self.attn( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + deterministic=deterministic, + init_cache=init_cache, + output_attentions=output_attentions, + ) + # residual connection + attn_output = attn_outputs[0] # output_attn: a, (attentions) + outputs = attn_outputs[1:] + # residual connection + hidden_states = attn_output + residual + # Cross-Attention Block + if encoder_hidden_states is not None: + # add one self-attention block for cross-attention + if not hasattr(self, "crossattention"): + raise ValueError( + f"If `encoder_hidden_states` are passed, {self} has to be instantiated with " + "cross-attention layers by setting `config.add_cross_attention=True`" + ) + residual = hidden_states + hidden_states = self.ln_cross_attn(hidden_states) + cross_attn_outputs = self.crossattention( + hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + position_ids=position_ids, + deterministic=deterministic, + output_attentions=output_attentions, + ) + attn_output = cross_attn_outputs[0] + # residual connection + hidden_states = residual + attn_output + outputs = outputs + cross_attn_outputs[1:] # add cross attentions if we output attention weights + residual = hidden_states + hidden_states = self.ln_2(hidden_states) + feed_forward_hidden_states = self.moe(hidden_states, deterministic=deterministic) + # residual connection + hidden_states = residual + feed_forward_hidden_states + outputs = (hidden_states,) + outputs + return outputs + + + +class LMCFlaxGPT2PreTrainedModel(FlaxPreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = GPT2Config + base_model_prefix = "transformer" + module_class: nn.Module = None + + def __init__( + self, + config: GPT2Config, + input_shape: Tuple = (1, 1), + seed: int = 0, + dtype: jnp.dtype = jnp.float32, + _do_init: bool = True, + **kwargs, + ): + module = self.module_class(config=config, dtype=dtype, **kwargs) + super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init) + + def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict: + # init input tensors + input_ids = jnp.zeros(input_shape, dtype="i4") + attention_mask = jnp.ones_like(input_ids) + position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape) + params_rng, dropout_rng = jax.random.split(rng) + rngs = {"params": params_rng, "dropout": dropout_rng} + + if self.config.add_cross_attention: + encoder_hidden_states = jnp.zeros(input_shape + (self.config.n_embd,)) + encoder_attention_mask = attention_mask + module_init_outputs = self.module.init( + rngs, + input_ids, + attention_mask, + position_ids, + encoder_hidden_states, + encoder_attention_mask, + return_dict=False, + ) + else: + module_init_outputs = self.module.init(rngs, input_ids, attention_mask, position_ids, return_dict=False) + + random_params = module_init_outputs["params"] + + if params is not None: + random_params = flatten_dict(unfreeze(random_params)) + params = flatten_dict(unfreeze(params)) + for missing_key in self._missing_keys: + params[missing_key] = random_params[missing_key] + self._missing_keys = set() + return freeze(unflatten_dict(params)) + else: + return random_params + + def init_cache(self, batch_size, max_length): + r""" + Args: + batch_size (`int`): + batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache. + max_length (`int`): + maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized + cache. + """ + # init input variables to retrieve cache + input_ids = jnp.ones((batch_size, max_length)) + attention_mask = jnp.ones_like(input_ids) + position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape) + + init_variables = self.module.init( + jax.random.PRNGKey(0), input_ids, attention_mask, position_ids, return_dict=False, init_cache=True + ) + return unfreeze(init_variables["cache"]) + + def __call__( + self, + input_ids, + attention_mask=None, + position_ids=None, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + params: Optional[dict] = None, + past_key_values: Optional[dict] = None, + dropout_rng: jax.random.PRNGKey = None, + train: bool = False, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ): + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if encoder_hidden_states is not None and encoder_attention_mask is None: + batch_size, sequence_length = encoder_hidden_states.shape[:2] + encoder_attention_mask = jnp.ones((batch_size, sequence_length)) + + batch_size, sequence_length = input_ids.shape + + if position_ids is None: + if past_key_values is not None: + raise ValueError("Make sure to provide `position_ids` when passing `past_key_values`.") + + position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)) + + if attention_mask is None: + attention_mask = jnp.ones((batch_size, sequence_length)) + + # Handle any PRNG if needed + rngs = {} + if dropout_rng is not None: + rngs["dropout"] = dropout_rng + + inputs = {"params": params or self.params} + + # if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that it can be changed by FlaxGPT2Attention module + if past_key_values: + inputs["cache"] = past_key_values + mutable = ["cache"] + else: + mutable = False + + outputs = self.module.apply( + inputs, + jnp.array(input_ids, dtype="i4"), + jnp.array(attention_mask, dtype="i4"), + jnp.array(position_ids, dtype="i4"), + encoder_hidden_states, + encoder_attention_mask, + not train, + False, + output_attentions, + output_hidden_states, + return_dict, + rngs=rngs, + mutable=mutable, + ) + + # add updated cache to model output + if past_key_values is not None and return_dict: + outputs, past_key_values = outputs + outputs["past_key_values"] = unfreeze(past_key_values["cache"]) + return outputs + elif past_key_values is not None and not return_dict: + outputs, past_key_values = outputs + outputs = outputs[:1] + (unfreeze(past_key_values["cache"]),) + outputs[1:] + + return outputs + + +class LMCFlaxGPT2BlockCollection(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.blocks = [ + LMCFlaxGPT2Block(self.config.lmc_config, name=str(i), dtype=self.dtype) if i in self.config.lmc_layer_indices + else LMCFlaxGPT2Block(self.config, name=str(i), dtype=self.dtype) + for i in range(self.config.num_hidden_layers) + ] + + def __call__( + self, + hidden_states, + attention_mask=None, + position_ids=None, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + deterministic: bool = True, + init_cache: bool = False, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ): + all_attentions = () if output_attentions else None + all_hidden_states = () if output_hidden_states else None + all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None + + for block in self.blocks: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = block( + hidden_states, + attention_mask, + position_ids=position_ids, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + deterministic=deterministic, + init_cache=init_cache, + output_attentions=output_attentions, + ) + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions += (layer_outputs[1],) + + if encoder_hidden_states is not None: + all_cross_attentions += (layer_outputs[2],) + + # this contains possible `None` values - `FlaxGPT2Module` will filter them out + outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions) + + return outputs + + +class LMCFlaxGPT2Module(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.embed_dim = self.config.hidden_size + + self.wte = nn.Embed( + self.config.vocab_size, + self.embed_dim, + embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range), + dtype=self.dtype, + ) + if self.config.position_embeddings == "learnable": + self.wpe = nn.Embed( + self.config.max_position_embeddings, + self.embed_dim, + embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range), + dtype=self.dtype, + ) + if self.config.position_embeddings == "sinusoidal": + self.wpe_table = create_sinusoidal_positions(self.config.max_position_embeddings, self.embed_dim) + self.dropout = nn.Dropout(rate=self.config.embd_pdrop) + self.h = LMCFlaxGPT2BlockCollection(self.config, dtype=self.dtype) + self.ln_f = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype) + + def __call__( + self, + input_ids, + attention_mask, + position_ids, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + deterministic=True, + init_cache: bool = False, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ): + input_embeds = self.wte(input_ids.astype("i4")) + if self.config.position_embeddings == "learnable": + position_embeds = self.wpe(position_ids.astype("i4")) + hidden_states = input_embeds + position_embeds + if self.config.position_embeddings == "sinusoidal": + position_embeds = jnp.take(self.wpe_table, position_ids, axis=0) + hidden_states = input_embeds + position_embeds + if self.config.position_embeddings == "rope": + hidden_states = input_embeds + hidden_states = self.dropout(hidden_states, deterministic=deterministic) + outputs = self.h( + hidden_states, + attention_mask, + position_ids=position_ids, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + deterministic=deterministic, + init_cache=init_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + hidden_states = self.ln_f(hidden_states) + + if output_hidden_states: + all_hidden_states = outputs[1] + (hidden_states,) + outputs = (hidden_states, all_hidden_states) + outputs[2:] + else: + outputs = (hidden_states,) + outputs[1:] + + if not return_dict: + return tuple(v for v in outputs if v is not None) + + return FlaxBaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + hidden_states=outputs[1], + attentions=outputs[2], + cross_attentions=outputs[3], + ) +class LMCFlaxGPT2Model(LMCFlaxGPT2PreTrainedModel): + module_class = LMCFlaxGPT2Module +class LMCFlaxGPT2LMHeadModule(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.transformer = LMCFlaxGPT2Module(self.config, dtype=self.dtype) + self.lm_head = nn.Dense( + self.config.vocab_size, + use_bias=False, + dtype=self.dtype, + kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range), + ) + def __call__( + self, + input_ids, + attention_mask, + position_ids, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + deterministic: bool = True, + init_cache: bool = False, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ): + outputs = self.transformer( + input_ids, + attention_mask, + position_ids, + encoder_hidden_states, + encoder_attention_mask, + deterministic=deterministic, + init_cache=init_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + + if self.config.tie_word_embeddings: + shared_kernel = self.transformer.variables["params"]["wte"]["embedding"].T + lm_logits = self.lm_head.apply({"params": {"kernel": shared_kernel}}, hidden_states) + else: + lm_logits = self.lm_head(hidden_states) + if not return_dict: + return (lm_logits,) + outputs[1:] + return FlaxCausalLMOutputWithCrossAttentions( + logits=lm_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + cross_attentions=outputs.cross_attentions, + ) + + + +class LMCFlaxGPT2LMHeadModel(LMCFlaxGPT2PreTrainedModel): + module_class = LMCFlaxGPT2LMHeadModule + + def prepare_inputs_for_generation(self, input_ids, max_length, attention_mask: Optional[jax.Array] = None): + # initializing the cache + batch_size, seq_length = input_ids.shape + + past_key_values = self.init_cache(batch_size, max_length) + # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length. + # But since GPT2 uses a causal mask, those positions are masked anyways. + # Thus we can create a single static attention_mask here, which is more efficient for compilation + extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4") + if attention_mask is not None: + position_ids = attention_mask.cumsum(axis=-1) - 1 + extended_attention_mask = lax.dynamic_update_slice( + extended_attention_mask, attention_mask.astype("i4"), (0, 0) + ) + else: + position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length)) + + return { + "past_key_values": past_key_values, + "attention_mask": extended_attention_mask, + "position_ids": position_ids, + } + + def update_inputs_for_generation(self, model_outputs, model_kwargs): + model_kwargs["past_key_values"] = model_outputs.past_key_values + model_kwargs["position_ids"] = model_kwargs["position_ids"][:, -1:] + 1 + return model_kwargs \ No newline at end of file diff --git a/src/lgmodeling/matching_utils.py b/src/lgmodeling/matching_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..073c0a0794ad7c3378f5eb26595e99b981747c27 --- /dev/null +++ b/src/lgmodeling/matching_utils.py @@ -0,0 +1,535 @@ +import os +import copy +import time +import random +import itertools +import numpy as np +import jax.lax as lax +import jax.numpy as jnp +from utils import rngmix +import matplotlib.pyplot as plt +from typing import NamedTuple +from collections import defaultdict +from flax.core import freeze, unfreeze +from scipy.optimize import linear_sum_assignment, minimize +from jax import random, tree_util, jit, grad, value_and_grad +def compute_objective(A, X, X_prime, Y, Y_prime): + A_inv = np.linalg.inv(A) + term1 = X - X_prime @ A.T + term2 = Y - Y_prime @ A_inv + return np.sum(term1**2) + np.sum(term2**2) + +def compute_gradient(A, X, X_prime, Y, Y_prime): + A_inv = np.linalg.inv(A) + term1 = -2 * X.T @ X_prime + 2 * A @ X_prime.T @ X_prime + term2 = 2 * A_inv.T @ Y_prime.T @ (Y - Y_prime @ A_inv) @ A_inv.T + return term1 + term2 + +def line_search(A, grad, X, X_prime, Y, Y_prime, max_step=1, tau=0.5, c1=1e-4): + eta = max_step + f_current = compute_objective(A, X, X_prime, Y, Y_prime) + grad_norm2 = np.sum(grad**2) + n = A.shape[0] + while eta > 1e-10: + A_new = A - eta * grad + if np.linalg.matrix_rank(A_new) < n: + eta *= tau + continue + f_new = compute_objective(A_new, X, X_prime, Y, Y_prime) + if f_new <= f_current - c1 * eta * grad_norm2: + return eta + eta *= tau + return 0 +@jit +def compute_objective_jax(A, X, X_prime, Y, Y_prime, cond_threshold=1e6): + cond = jnp.linalg.cond(A) + def safe_obj(): + A_inv = jnp.linalg.inv(A) + term1 = X - X_prime @ A.T + term2 = Y - Y_prime @ A_inv + return jnp.sum(term1**2) + jnp.sum(term2**2) + return lax.cond(cond > cond_threshold, lambda: jnp.inf, safe_obj) +compute_value_and_grad_jax = jit(value_and_grad(compute_objective_jax)) +def solve_orthogonal(X, X_prime, Y, Y_prime): + B = X.T @ X_prime + Y.T @ Y_prime + U, _, Vt = np.linalg.svd(B) + return U @ Vt +def solve_rope(X, X_prime, Y, Y_prime,max_iters=200, tol=1e-16): + d = X.shape[1] + assert d % 2 == 0, "d must be even." + assert X.shape[1] == X_prime.shape[1] == Y.shape[1] == Y_prime.shape[1] + def rot(theta): + c, s = np.cos(theta), np.sin(theta) + return np.array([[c, -s], [s, c]]) + def block_cols(j): return [2*j, 2*j+1] + def solve_block(Q1blk, Q2blk, K1blk, K2blk): + A, B, Ah, Bh = Q1blk, Q2blk, K1blk, K2blk + a, ah = np.sum(A*A), np.sum(Ah*Ah) + c_const = np.sum(B*B) + np.sum(Bh*Bh) + C, Ch = A.T @ B, Ah.T @ Bh + t_tr, s_sk = np.trace(C), C[0,1] - C[1,0] + th_tr, sh_sk = np.trace(Ch), Ch[0,1] - Ch[1,0] + u, v, w = t_tr**2+s_sk**2, th_tr**2+sh_sk**2, t_tr*th_tr+s_sk*sh_sk + eps = 1e-18 + def phi(t): return max(u*t + v/max(t,eps) + 2*w, 0.0) + def gprime(t): + denom = np.sqrt(phi(t)) + if denom < eps: return a - ah/(t*t) + return (a - ah/(t*t)) - (u - v/(t*t)) / denom + t0 = np.sqrt((ah+eps)/(a+eps)) + t_lo, gp_lo = t0, gprime(t0) + if gp_lo < 0.0: + t_hi = t_lo + for _ in range(max_iters): + t_hi *= 2.0 + if gprime(t_hi) >= 0.0: break + else: + t_hi = t_lo + for _ in range(max_iters): + t_lo *= 0.5 + if gprime(t_lo) <= 0.0: break + def gval(t): return a*t + ah/max(t,eps) + c_const - 2*np.sqrt(phi(t)) + if not (gprime(t_lo) <= 0.0 <= gprime(t_hi)): + t_star = min([(t_lo,gval(t_lo)),(t_hi,gval(t_hi))], key=lambda z:z[1])[0] + else: + for _ in range(max_iters): + t_mid = 0.5*(t_lo+t_hi) + gp_mid = gprime(t_mid) + if abs(gp_mid) < tol or (t_hi-t_lo) <= tol*(1+t_mid): + t_star = t_mid; break + if gp_mid < 0.0: t_lo = t_mid + else: t_hi = t_mid + else: + t_star = 0.5*(t_lo+t_hi) + rho = np.sqrt(max(t_star, eps)) + alpha = rho*t_tr + (1/rho)*th_tr + beta = rho*s_sk + (1/rho)*sh_sk + theta = np.arctan2(beta, alpha) + return rho, theta + P = np.zeros((d, d)) + for j in range(d//2): + cols = block_cols(j) + rho, theta = solve_block(X[:,cols], X_prime[:,cols], Y[:,cols], Y_prime[:,cols]) + P[np.ix_(cols, cols)] = rho * rot(theta) + return P + + +def optimize_alignment(A_init, X, X_prime, Y, Y_prime, max_iter=5000): + objective_values = [] + grad_norms = [] + condition_nums = [] + + def obj_fn(flat_A): + A = flat_A.reshape(A_init.shape) + obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime)) + return float(obj), np.array(grad_val).flatten() + + def callback(flat_A): + A = flat_A.reshape(A_init.shape) + obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime)) + grad_norm = jnp.linalg.norm(grad_val, 'fro') + cond = jnp.linalg.cond(jnp.array(A)) + objective_values.append(float(obj)) + grad_norms.append(float(grad_norm)) + condition_nums.append(float(cond)) + + res = minimize(obj_fn, A_init.flatten(), jac=True, method='L-BFGS-B', options={'maxiter': max_iter}, callback=callback) + A_opt = res.x.reshape(A_init.shape) + return A_opt, objective_values, grad_norms, condition_nums + + +def extract_attention_params(attn): + c_attn_kernel = np.array(attn['c_attn']['kernel']) + c_attn_bias = np.array(attn['c_attn']['bias']) + c_proj_kernel = np.array(attn['c_proj']['kernel']) + c_proj_bias = np.array(attn['c_proj']['bias']) + query, key, value = np.split(c_attn_kernel, 3, axis=0) + query_bias, key_bias, value_bias = np.split(c_attn_bias, 3, axis=0) + return query, key, value, query_bias, key_bias, value_bias, c_proj_kernel, c_proj_bias +def reshape_attention_weights(query, key, value, query_bias, key_bias, value_bias, out_kernel, num_heads): + D = query.shape[1] + D_k = D_v = D // num_heads + def stack_per_head(tensor, axis=0): + return np.stack([ + tensor[i * D_k:(i + 1) * D_k, :].T if axis == 0 else tensor[:, i * D_k:(i + 1) * D_k].T + for i in range(num_heads) + ]) + def stack_bias_per_head(bias): + return np.stack([bias[i * D_k:(i + 1) * D_k].T for i in range(num_heads)]) + W_Q = stack_per_head(query) + W_K = stack_per_head(key) + W_V = stack_per_head(value) + W_O = stack_per_head(out_kernel, axis=1) + b_Q = stack_bias_per_head(query_bias) + b_K = stack_bias_per_head(key_bias) + b_V = stack_bias_per_head(value_bias) + return W_Q, b_Q, W_K, b_K, W_V, b_V, W_O +def compute_extended_weights(W, b): + return np.vstack([W, b.reshape(1, -1)]) +def compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + h, activations, alpha=0.5): + C = np.zeros((h, h)) + for i in range(h): + tilde_W_Q_a_i = np.vstack([W_Q_a[i], b_Q_a[i].reshape(1, -1)]) + tilde_W_K_a_i = np.vstack([W_K_a[i], b_K_a[i].reshape(1, -1)]) + tilde_W_V_a_i = np.vstack([W_V_a[i], b_V_a[i].reshape(1, -1)]) + QKT_a_i = tilde_W_Q_a_i @ tilde_W_K_a_i.T + VO_a_i = tilde_W_V_a_i @ W_O_a[i] + centered_QKT_a_i = QKT_a_i - np.mean(QKT_a_i, axis=1, keepdims=True) + + for j in range(h): + tilde_W_Q_b_j = np.vstack([W_Q_b[j], b_Q_b[j].reshape(1, -1)]) + tilde_W_K_b_j = np.vstack([W_K_b[j], b_K_b[j].reshape(1, -1)]) + tilde_W_V_b_j = np.vstack([W_V_b[j], b_V_b[j].reshape(1, -1)]) + QKT_b_j = tilde_W_Q_b_j @ tilde_W_K_b_j.T + VO_b_j = tilde_W_V_b_j @ W_O_b[j] + centered_QKT_b_j = QKT_b_j - np.mean(QKT_b_j, axis=1, keepdims=True) + + cost = 0.5 * np.sum((centered_QKT_a_i - centered_QKT_b_j) ** 2) + cost += 0.5 * np.sum((VO_a_i - VO_b_j) ** 2) + C[i, j] = cost + return C +def additive_align_single_head(W_Q_a_i, b_Q_a_i, W_K_a_i, b_K_a_i, W_V_a_i, b_V_a_i, W_O_a_i, + W_Q_b_i, b_Q_b_i, W_K_b_i, b_K_b_i, W_V_b_i, b_V_b_i, W_O_b_i, optimize): + tilde_W_Q_a_i = compute_extended_weights(W_Q_a_i, b_Q_a_i) + tilde_W_K_a_i = compute_extended_weights(W_K_a_i, b_K_a_i) + tilde_W_V_a_i = compute_extended_weights(W_V_a_i, b_V_a_i) + Y_O_a_i = W_O_a_i.T + tilde_W_Q_b_i = compute_extended_weights(W_Q_b_i, b_Q_b_i) + tilde_W_K_b_i = compute_extended_weights(W_K_b_i, b_K_b_i) + tilde_W_V_b_i = compute_extended_weights(W_V_b_i, b_V_b_i) + Y_O_b_i = W_O_b_i.T + A_init = solve_orthogonal(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i) + B_init = solve_orthogonal(Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i) + if optimize: + A, objective_values_A, grad_norms_A, condition_nums_A = optimize_alignment( + A_init, tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i + ) + B, objective_values_B, grad_norms_B, condition_nums_B = optimize_alignment( + B_init, Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i + ) + else: + A = A_init + B = B_init + A_inv = np.linalg.inv(A) + B_inv = np.linalg.inv(B) + W_Q_aligned = W_Q_b_i @ A.T + b_Q_aligned = b_Q_b_i @ A.T + W_K_aligned = W_K_b_i @ A_inv + b_K_aligned = b_K_b_i @ A_inv + W_V_aligned = W_V_b_i @ B_inv + b_V_aligned = b_V_b_i @ B_inv + W_O_aligned = B @ W_O_b_i + aligned_params = { + 'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned}, + 'key': {'kernel': W_K_aligned, 'bias': b_K_aligned}, + 'value': {'kernel': W_V_aligned, 'bias': b_V_aligned}, + 'out': {'kernel': W_O_aligned} + } + if optimize: + return { + 'aligned_params': aligned_params, + 'metrics_A': { + 'objective_values': objective_values_A, + 'grad_norms': grad_norms_A, + 'condition_nums': condition_nums_A + }, + 'metrics_B': { + 'objective_values': objective_values_B, + 'grad_norms': grad_norms_B, + 'condition_nums': condition_nums_B + } + } + return {'aligned_params': aligned_params} + +def _fro2(x): + if x.ndim == 1: # vector -> Euclidean norm + return float(np.linalg.norm(x)**2) + else: # matrix -> Frobenius norm + return float(np.linalg.norm(x, 'fro')**2) + +def rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, h): + aligned_params = {} + + total_pre = 0.0 + total_post = 0.0 + # --- Pairwise totals --- + total_qk_pre = 0.0 + total_qk_post = 0.0 + total_vo_pre = 0.0 + total_vo_post = 0.0 + + for i in range(h): + # ===== Build augmented (kernel+bias row) for A (Q,K) ===== + tilde_W_Q_a_i = np.vstack([W_Q_a[i], b_Q_a[i].reshape(1, -1)]) + tilde_W_K_a_i = np.vstack([W_K_a[i], b_K_a[i].reshape(1, -1)]) + + # ===== Build augmented (kernel+bias row) for B (Q,K) ===== + tilde_W_Q_b_i = np.vstack([W_Q_b[i], b_Q_b[i].reshape(1, -1)]) + tilde_W_K_b_i = np.vstack([W_K_b[i], b_K_b[i].reshape(1, -1)]) + + # ===== Solve A_i for (Q,K) pair ===== + A_i_init = solve_rope(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i) + A_i = A_i_init # (optionally run a refinement step) + + # ===== Build augmented (kernel+bias row) for (V,O) pair ===== + tilde_W_V_a_i = np.vstack([W_V_a[i], b_V_a[i].reshape(1, -1)]) + tilde_W_V_b_i = np.vstack([W_V_b[i], b_V_b[i].reshape(1, -1)]) + + # For O we transpose and pad with one zero row to match augmented shape + Y_O = W_O_a[i].T + Y_O_prime = W_O_b[i].T + Y_O_padded = np.vstack([Y_O, np.zeros((1, Y_O.shape[1]))]) + Y_O_prime_padded = np.vstack([Y_O_prime, np.zeros((1, Y_O_prime.shape[1]))]) + + # ===== Solve B_i for (V,O) pair ===== + B_i_init = solve_rope(tilde_W_V_a_i, tilde_W_V_b_i, Y_O_padded, Y_O_prime_padded) + B_i = B_i_init # (optionally run a refinement step) + + # ===== Apply transforms ===== + A_i_inv = np.linalg.inv(A_i) + B_i_inv = np.linalg.inv(B_i) + + # W_Q_aligned = W_Q_b[i] @ A_i.T + # b_Q_aligned = b_Q_b[i] @ A_i.T + + # W_K_aligned = W_K_b[i] @ A_i_inv + # b_K_aligned = b_K_b[i] @ A_i_inv + + # W_V_aligned = W_V_b[i] @ B_i_inv + # b_V_aligned = b_V_b[i] @ B_i_inv + + # W_O_aligned = B_i @ W_O_b[i] + W_Q_aligned = W_Q_b[i] @ A_i_inv.T # was A_i.T -> FIX: A_i^{-T} + b_Q_aligned = b_Q_b[i] @ A_i_inv.T + + W_K_aligned = W_K_b[i] @ A_i # was A_i_inv -> FIX: A_i + b_K_aligned = b_K_b[i] @ A_i + + # --- V,O pair: use P for V (right-multiply), and P^{-1} for O (left-multiply) --- + W_V_aligned = W_V_b[i] @ B_i # was B_i_inv -> FIX: B_i + b_V_aligned = b_V_b[i] @ B_i + + W_O_aligned = B_i_inv @ W_O_b[i] # was B_i @ W_O_b[i] -> FIX: B_i^{-1} on the left + aligned_params[f'head_{i}'] = { + 'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned}, + 'key': {'kernel': W_K_aligned, 'bias': b_K_aligned}, + 'value': {'kernel': W_V_aligned, 'bias': b_V_aligned}, + 'out': {'kernel': W_O_aligned} + } + + # ===== Frobenius^2 BEFORE (a vs raw b) ===== + pre_q = _fro2(W_Q_a[i] - W_Q_b[i]) + _fro2(b_Q_a[i] - b_Q_b[i]) + pre_k = _fro2(W_K_a[i] - W_K_b[i]) + _fro2(b_K_a[i] - b_K_b[i]) + pre_v = _fro2(W_V_a[i] - W_V_b[i]) + _fro2(b_V_a[i] - b_V_b[i]) + pre_o = _fro2(W_O_a[i] - W_O_b[i]) # O has no bias in your structure + + # ===== Frobenius^2 AFTER (a vs aligned b) ===== + post_q = _fro2(W_Q_a[i] - W_Q_aligned) + _fro2(b_Q_a[i] - b_Q_aligned) + post_k = _fro2(W_K_a[i] - W_K_aligned) + _fro2(b_K_a[i] - b_K_aligned) + post_v = _fro2(W_V_a[i] - W_V_aligned) + _fro2(b_V_a[i] - b_V_aligned) + post_o = _fro2(W_O_a[i] - W_O_aligned) + + pre_sum = pre_q + pre_k + pre_v + pre_o + post_sum = post_q + post_k + post_v + post_o + + total_pre += pre_sum + total_post += post_sum + + # ===== Pairwise sums ===== + pre_qk = pre_q + pre_k + post_qk = post_q + post_k + pre_vo = pre_v + pre_o + post_vo = post_v + post_o + + total_qk_pre += pre_qk + total_qk_post += post_qk + total_vo_pre += pre_vo + total_vo_post += post_vo + + # ===== Per-head print ===== + print(f"[Head {i}] Fro^2 pre={pre_sum:.6f} post={post_sum:.6f} improve={pre_sum - post_sum:.6f}") + print(f" Q: pre={pre_q:.6f} post={post_q:.6f}") + print(f" K: pre={pre_k:.6f} post={post_k:.6f}") + print(f" V: pre={pre_v:.6f} post={post_v:.6f}") + print(f" O: pre={pre_o:.6f} post={post_o:.6f}") + # --- New: pairwise breakdowns --- + print(f" [Q,K] pair: pre={pre_qk:.6f} post={post_qk:.6f} improve={pre_qk - post_qk:.6f}") + print(f" [V,O] pair: pre={pre_vo:.6f} post={post_vo:.6f} improve={pre_vo - post_vo:.6f}") + # ===== Totals ===== + print("=== Frobenius^2 (including biases) ===") + print(f"Total pre : {total_pre:.6f}") + print(f"Total post: {total_post:.6f}") + print(f"Total improvement: {total_pre - total_post:.6f} ({0.0 if total_pre==0 else 100.0*(total_pre-total_post)/total_pre:.2f}%)") + + # --- New: Pairwise totals --- + print("=== Pairwise Frobenius^2 (including biases) ===") + print(f"[Q,K] total pre : {total_qk_pre:.6f}") + print(f"[Q,K] total post: {total_qk_post:.6f}") + print(f"[Q,K] improvement: {total_qk_pre - total_qk_post:.6f} ({0.0 if total_qk_pre==0 else 100.0*(total_qk_pre-total_qk_post)/total_qk_pre:.2f}%)") + print(f"[V, O] total pre : {total_vo_pre:.6f}") + print(f"[V, O] total post: {total_vo_post:.6f}") + print(f"[V, O] improvement: {total_vo_pre - total_vo_post:.6f} ({0.0 if total_vo_pre==0 else 100.0*(total_vo_pre-total_vo_post)/total_vo_pre:.2f}%)") + + return aligned_params +def merge_aligned_params(aligned_params, h, D, out_bias_b): + query_kernel = np.stack([aligned_params[f'head_{i}']['query']['kernel'] for i in range(h)], axis=1) + query_bias = np.stack([aligned_params[f'head_{i}']['query']['bias'] for i in range(h)], axis=0) + key_kernel = np.stack([aligned_params[f'head_{i}']['key']['kernel'] for i in range(h)], axis=1) + key_bias = np.stack([aligned_params[f'head_{i}']['key']['bias'] for i in range(h)], axis=0) + value_kernel = np.stack([aligned_params[f'head_{i}']['value']['kernel'] for i in range(h)], axis=1) + value_bias = np.stack([aligned_params[f'head_{i}']['value']['bias'] for i in range(h)], axis=0) + out_kernel = np.stack([aligned_params[f'head_{i}']['out']['kernel'] for i in range(h)], axis=0) + + query_kernel = query_kernel.transpose(1, 2, 0).reshape(-1, D) + key_kernel = key_kernel.transpose(1, 2, 0).reshape(-1, D) + value_kernel = value_kernel.transpose(1, 2, 0).reshape(-1, D) + out_kernel = out_kernel.transpose(2, 0, 1).reshape(D, -1) + + query_bias = query_bias.reshape(-1) + key_bias = key_bias.reshape(-1) + value_bias = value_bias.reshape(-1) + + return { + 'c_attn': { + 'kernel': jnp.array(np.concatenate([query_kernel, key_kernel, value_kernel], axis=0)), + 'bias': jnp.array(np.concatenate([query_bias, key_bias, value_bias], axis=0)), + }, + 'c_proj': {'kernel': jnp.array(out_kernel),'bias': jnp.array(out_bias_b),} + } + +def align_attention_params(rng, params_a, params_b, layer_idx, config, activation, permute_heads=True, optimize=False, alpha=0.5): + num_heads = config.lmc_config.n_head + attn_a = params_a['transformer']['h'][str(layer_idx)]['attn'] + attn_b = params_b['transformer']['h'][str(layer_idx)]['attn'] + query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, out_bias_a = extract_attention_params(attn_a) + query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, out_bias_b = extract_attention_params(attn_b) + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_attention_weights(query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, num_heads) + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_attention_weights(query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, num_heads) + if permute_heads: + C = compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, num_heads, activation, alpha) + row_ind, col_ind = linear_sum_assignment(C) + print("Best Permutation Heads:", col_ind) + W_Q_b = [W_Q_b[j] for j in col_ind] + b_Q_b = [b_Q_b[j] for j in col_ind] + W_K_b = [W_K_b[j] for j in col_ind] + b_K_b = [b_K_b[j] for j in col_ind] + W_V_b = [W_V_b[j] for j in col_ind] + b_V_b = [b_V_b[j] for j in col_ind] + W_O_b = [W_O_b[j] for j in col_ind] + if optimize: + metrics_A_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']} + metrics_B_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']} + aligned_params, return_dict = {}, {} + if(config.position_embeddings in ["learnable","sinusoidal"]): + for i in range(num_heads): + result = additive_align_single_head( + W_Q_a[i], b_Q_a[i], W_K_a[i], b_K_a[i], W_V_a[i], b_V_a[i], W_O_a[i], + W_Q_b[i], b_Q_b[i], W_K_b[i], b_K_b[i], W_V_b[i], b_V_b[i], W_O_b[i], optimize + ) + aligned_params[f'head_{i}'] = result['aligned_params'] + if optimize: + for key in metrics_A_all: + metrics_A_all[key].append(result['metrics_A'][key]) + metrics_B_all[key].append(result['metrics_B'][key]) + return_dict['aligned_params'] = merge_aligned_params(aligned_params, num_heads, query_a.shape[1], out_bias_b) + if optimize: + return_dict['metrics_A_all'] = metrics_A_all + return_dict['metrics_B_all'] = metrics_B_all + elif(config.position_embeddings in ["rope"]): + aligned_params = rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads) + return return_dict +def permute_align_attention_params(rng, params_a, params_b, layer_idx, config,col_ind): + num_heads = config.lmc_config.n_head + attn_a = params_a['transformer']['h'][layer_idx]['attn'] + attn_b = params_b['transformer']['h'][layer_idx]['attn'] + query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, out_bias_a = extract_attention_params(attn_a) + query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, out_bias_b = extract_attention_params(attn_b) + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_attention_weights(query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, num_heads) + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_attention_weights(query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, num_heads) + W_Q_b = [W_Q_b[j] for j in col_ind] + b_Q_b = [b_Q_b[j] for j in col_ind] + W_K_b = [W_K_b[j] for j in col_ind] + b_K_b = [b_K_b[j] for j in col_ind] + W_V_b = [W_V_b[j] for j in col_ind] + b_V_b = [b_V_b[j] for j in col_ind] + W_O_b = [W_O_b[j] for j in col_ind] + if(config.position_embeddings in ["learnable","sinusoidal"]): + aligned_params = additive_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads) + elif(config.position_embeddings in ["rope"]): + aligned_params = rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads) + return merge_aligned_params(aligned_params, num_heads, query_a.shape[1], out_bias_b) +def naive_align_attention_params(rng, params_a, params_b, layer_idx, config): + num_heads = config.lmc_config.n_head + attn_a = params_a['transformer']['h'][layer_idx]['attn'] + attn_b = params_b['transformer']['h'][layer_idx]['attn'] + query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, out_bias_a = extract_attention_params(attn_a) + query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, out_bias_b = extract_attention_params(attn_b) + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_attention_weights(query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, num_heads) + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_attention_weights(query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, num_heads) + if(config.position_embeddings in ["learnable","sinusoidal"]): + aligned_params = additive_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads) + elif(config.position_embeddings in ["rope"]): + aligned_params = rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads) + return merge_aligned_params(aligned_params, num_heads, query_a.shape[1], out_bias_b) +def all_matching_attn(rng, params_a, params_b, config): + results = {} + permutations = list(itertools.permutations(range(config.lmc_config.n_head))) + if config.lmc_config.n_head > 4: + permutations = random.sample(permutations, 24) + for perm in permutations: + print("Permutation",perm) + temp_params = copy.deepcopy(params_b) + for layer_idx in config.lmc_layer_indices: + aligned_attention_params = permute_align_attention_params(rng, params_a, params_b, str(layer_idx), config, perm) + temp_params['transformer']['h'][str(layer_idx)]['attn'] = aligned_attention_params + results[str(perm)] = temp_params + return results + +def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ + ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), + ] + for name, init_method, permute_heads, optimize in configurations: + aligned_params = copy.deepcopy(params_b) + if optimize: + layer_to_metrics_A = {} + layer_to_metrics_B = {} + for layer_idx in config.lmc_layer_indices: + if activation is not None: activations_for_layer = activation[layer_idx] + else: activations_for_layer = None + result = align_attention_params( + rng, params_a, aligned_params, layer_idx, config, + activations_for_layer, permute_heads=permute_heads, optimize=optimize + ) + aligned_params['transformer']['h'][str(layer_idx)]['attn'] = result['aligned_params'] + if optimize: + layer_to_metrics_A[layer_idx] = result['metrics_A_all'] + layer_to_metrics_B[layer_idx] = result['metrics_B_all'] + + total_sum = tree_util.tree_reduce(lambda acc, x: acc + jnp.sum(x), aligned_params, initializer=0) + print(f"{name}: {total_sum}, sanity check") + params_dict[name] = aligned_params + return params_dict + + # cost_head = copy.deepcopy(params_b) + # naive_head = copy.deepcopy(params_b) + # for layer_idx in config.lmc_layer_indices: + # aligned_attention_params = cost_align_attention_params(rng, params_a, params_b, str(layer_idx), config) + # cost_head['transformer']['h'][str(layer_idx)]['attn'] = aligned_attention_params + # for layer_idx in config.lmc_layer_indices: + # aligned_attention_params = naive_align_attention_params(rng, params_a, params_b, str(layer_idx), config) + # naive_head['transformer']['h'][str(layer_idx)]['attn'] = aligned_attention_params + # return {"cost_head": cost_head, "naive_head": naive_head} diff --git a/src/lgmodeling/model.py b/src/lgmodeling/model.py new file mode 100644 index 0000000000000000000000000000000000000000..e09fc265904ef070e0c61d901bf7c832b9681159 --- /dev/null +++ b/src/lgmodeling/model.py @@ -0,0 +1,601 @@ +from typing import Optional, Tuple, List +import flax.linen as nn +import jax +import jax.numpy as jnp +from jax import lax +from flax.core.frozen_dict import FrozenDict, freeze, unfreeze +from flax.linen.attention import dot_product_attention_weights +from flax.traverse_util import flatten_dict, unflatten_dict +from flax.linen import combine_masks, make_causal_mask +from transformers.models.gpt2.modeling_flax_gpt2 import ( + FlaxConv1D, + FlaxGPT2MLP, + FlaxGPT2Block, + GPT2Config, + FlaxBaseModelOutputWithPastAndCrossAttentions, + FlaxCausalLMOutputWithCrossAttentions, +) +from transformers.modeling_flax_utils import ACT2FN, FlaxPreTrainedModel +from typing import Callable, Literal +def print_model(flax_params, file=None): + flat_params = flatten_dict(flax_params) + for path, value in flat_params.items(): + name = "/".join(path) + line = f"{name} {value.shape}" + if file: + print(line, file=file) + else: + print(line) +class LMCFlaxGPT2Attention(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + causal: bool = True + is_cross_attention: bool = False + def setup(self): + config = self.config + self.embed_dim = config.hidden_size + self.num_heads = config.att_num_heads + self.head_dim = self.embed_dim // self.num_heads + if self.is_cross_attention: + self.c_attn = FlaxConv1D(2 * self.embed_dim, dtype=self.dtype) + self.q_attn = FlaxConv1D(self.embed_dim, dtype=self.dtype) + else: + self.c_attn = FlaxConv1D(3 * self.embed_dim, dtype=self.dtype) + self.c_proj = FlaxConv1D(self.embed_dim, dtype=self.dtype) + self.resid_dropout = nn.Dropout(rate=config.resid_pdrop) + if self.causal: + self.causal_mask = make_causal_mask( + jnp.ones((1, config.max_position_embeddings), dtype="bool"), dtype="bool" + ) + def _split_heads(self, hidden_states): + return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads, self.head_dim)) + def _merge_heads(self, hidden_states): + return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,)) + @nn.compact + def _concatenate_to_cache(self, key, value, query, attention_mask): + """ + This function takes projected key, value states from a single input token and concatenates the states to cached + states from previous steps. This function is slightly adapted from the official Flax repository: + https://github.com/google/flax/blob/491ce18759622506588784b4fca0e4bf05f8c8cd/flax/linen/attention.py#L252 + """ + # detect if we're initializing by absence of existing cache data. + is_initialized = self.has_variable("cache", "cached_key") + cached_key = self.variable("cache", "cached_key", jnp.zeros, key.shape, key.dtype) + cached_value = self.variable("cache", "cached_value", jnp.zeros, value.shape, value.dtype) + cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32)) + if is_initialized: + *batch_dims, max_length, num_heads, depth_per_head = cached_key.value.shape + # update key, value caches with our new 1d spatial slices + cur_index = cache_index.value + indices = (0,) * len(batch_dims) + (cur_index, 0, 0) + key = lax.dynamic_update_slice(cached_key.value, key, indices) + value = lax.dynamic_update_slice(cached_value.value, value, indices) + cached_key.value = key + cached_value.value = value + num_updated_cache_vectors = query.shape[1] + cache_index.value = cache_index.value + num_updated_cache_vectors + # causal mask for cached decoder self-attention: our single query position should only attend to those key positions that have already been generated and cached, not the remaining zero elements. + pad_mask = jnp.broadcast_to( + jnp.arange(max_length) < cur_index + num_updated_cache_vectors, + tuple(batch_dims) + (1, num_updated_cache_vectors, max_length), + ) + attention_mask = combine_masks(pad_mask, attention_mask) + return key, value, attention_mask + + def __call__( + self, + hidden_states, + key_value_states: Optional[jnp.ndarray] = None, + attention_mask=None, + deterministic: bool = True, + init_cache: bool = False, + output_attentions: bool = False, + ): + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + batch_size = hidden_states.shape[0] + if not is_cross_attention: + qkv_out = self.c_attn(hidden_states) + query, key, value = jnp.split(qkv_out, 3, axis=2) + else: + q_out = self.q_attn(hidden_states) + (query,) = jnp.split(q_out, 1, axis=2) + kv_out = self.c_attn(key_value_states) + key, value = jnp.split(kv_out, 2, axis=2) + query = self._split_heads(query) + key = self._split_heads(key) + value = self._split_heads(value) + query_length, key_length = query.shape[1], key.shape[1] + if self.causal: + if self.has_variable("cache", "cached_key"): + mask_shift = self.variables["cache"]["cache_index"] + max_decoder_length = self.variables["cache"]["cached_key"].shape[1] + causal_mask = lax.dynamic_slice( + self.causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length) + ) + else: + causal_mask = self.causal_mask[:, :, :query_length, :key_length] + causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:]) + # combine masks if needed + if attention_mask is not None and self.causal: + attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape) + attention_mask = combine_masks(attention_mask, causal_mask) + elif self.causal: + attention_mask = causal_mask + elif attention_mask is not None: + attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2)) + dropout_rng = None + if not deterministic and self.config.attn_pdrop > 0.0: + dropout_rng = self.make_rng("dropout") + # During fast autoregressive decoding, we feed one position at a time, + # and cache the keys and values step by step. + if self.causal and (self.has_variable("cache", "cached_key") or init_cache): + key, value, attention_mask = self._concatenate_to_cache(key, value, query, attention_mask) + # transform boolean mask into float mask + if attention_mask is not None: + attention_bias = lax.select( + attention_mask > 0, + jnp.full(attention_mask.shape, 0.0).astype(self.dtype), + jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype), + ) + else: + attention_bias = None + # usual dot product attention + attn_weights = dot_product_attention_weights( + query, key, + bias=attention_bias, + dropout_rng=dropout_rng, + dropout_rate=self.config.attn_pdrop, + deterministic=deterministic, + dtype=self.dtype, + precision=None, + ) + attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value) + attn_output = self._merge_heads(attn_output) + attn_output = self.c_proj(attn_output) + attn_output = self.resid_dropout(attn_output, deterministic=deterministic) + outputs = (attn_output, attn_weights) if output_attentions else (attn_output,) + return outputs +class LMCFlaxGPT2Block(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + def setup(self): + hidden_size = self.config.hidden_size + inner_dim = self.config.n_inner if self.config.n_inner is not None else 4 * hidden_size + self.ln_1 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype) + self.attn = LMCFlaxGPT2Attention(self.config, dtype=self.dtype) + self.ln_2 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype) + if self.config.add_cross_attention: + self.crossattention = LMCFlaxGPT2Attention( + config=self.config, dtype=self.dtype, causal=False, is_cross_attention=True + ) + self.ln_cross_attn = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype) + self.mlp = FlaxGPT2MLP(self.config, inner_dim, dtype=self.dtype) + + def __call__( + self, + hidden_states, + attention_mask=None, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + deterministic: bool = True, + init_cache: bool = False, + output_attentions: bool = False, + ): + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + attn_outputs = self.attn( + hidden_states, + attention_mask=attention_mask, + deterministic=deterministic, + init_cache=init_cache, + output_attentions=output_attentions, + ) + # residual connection + attn_output = attn_outputs[0] # output_attn: a, (attentions) + outputs = attn_outputs[1:] + # residual connection + hidden_states = attn_output + residual + # Cross-Attention Block + if encoder_hidden_states is not None: + # add one self-attention block for cross-attention + if not hasattr(self, "crossattention"): + raise ValueError( + f"If `encoder_hidden_states` are passed, {self} has to be instantiated with " + "cross-attention layers by setting `config.add_cross_attention=True`" + ) + residual = hidden_states + hidden_states = self.ln_cross_attn(hidden_states) + cross_attn_outputs = self.crossattention( + hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + deterministic=deterministic, + output_attentions=output_attentions, + ) + attn_output = cross_attn_outputs[0] + # residual connection + hidden_states = residual + attn_output + outputs = outputs + cross_attn_outputs[1:] # add cross attentions if we output attention weights + residual = hidden_states + hidden_states = self.ln_2(hidden_states) + feed_forward_hidden_states = self.mlp(hidden_states, deterministic=deterministic) + # residual connection + hidden_states = residual + feed_forward_hidden_states + outputs = (hidden_states,) + outputs + return outputs + + + +class LMCFlaxGPT2PreTrainedModel(FlaxPreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = GPT2Config + base_model_prefix = "transformer" + module_class: nn.Module = None + + def __init__( + self, + config: GPT2Config, + input_shape: Tuple = (1, 1), + seed: int = 0, + dtype: jnp.dtype = jnp.float32, + _do_init: bool = True, + **kwargs, + ): + module = self.module_class(config=config, dtype=dtype, **kwargs) + super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init) + + def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict: + # init input tensors + input_ids = jnp.zeros(input_shape, dtype="i4") + attention_mask = jnp.ones_like(input_ids) + position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape) + params_rng, dropout_rng = jax.random.split(rng) + rngs = {"params": params_rng, "dropout": dropout_rng} + + if self.config.add_cross_attention: + encoder_hidden_states = jnp.zeros(input_shape + (self.config.n_embd,)) + encoder_attention_mask = attention_mask + module_init_outputs = self.module.init( + rngs, + input_ids, + attention_mask, + position_ids, + encoder_hidden_states, + encoder_attention_mask, + return_dict=False, + ) + else: + module_init_outputs = self.module.init(rngs, input_ids, attention_mask, position_ids, return_dict=False) + + random_params = module_init_outputs["params"] + + if params is not None: + random_params = flatten_dict(unfreeze(random_params)) + params = flatten_dict(unfreeze(params)) + for missing_key in self._missing_keys: + params[missing_key] = random_params[missing_key] + self._missing_keys = set() + return freeze(unflatten_dict(params)) + else: + return random_params + + def init_cache(self, batch_size, max_length): + r""" + Args: + batch_size (`int`): + batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache. + max_length (`int`): + maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized + cache. + """ + # init input variables to retrieve cache + input_ids = jnp.ones((batch_size, max_length)) + attention_mask = jnp.ones_like(input_ids) + position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape) + + init_variables = self.module.init( + jax.random.PRNGKey(0), input_ids, attention_mask, position_ids, return_dict=False, init_cache=True + ) + return unfreeze(init_variables["cache"]) + + def __call__( + self, + input_ids, + attention_mask=None, + position_ids=None, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + params: Optional[dict] = None, + past_key_values: Optional[dict] = None, + dropout_rng: jax.random.PRNGKey = None, + train: bool = False, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ): + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if encoder_hidden_states is not None and encoder_attention_mask is None: + batch_size, sequence_length = encoder_hidden_states.shape[:2] + encoder_attention_mask = jnp.ones((batch_size, sequence_length)) + + batch_size, sequence_length = input_ids.shape + + if position_ids is None: + if past_key_values is not None: + raise ValueError("Make sure to provide `position_ids` when passing `past_key_values`.") + + position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)) + + if attention_mask is None: + attention_mask = jnp.ones((batch_size, sequence_length)) + + # Handle any PRNG if needed + rngs = {} + if dropout_rng is not None: + rngs["dropout"] = dropout_rng + + inputs = {"params": params or self.params} + + # if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that it can be changed by FlaxGPT2Attention module + if past_key_values: + inputs["cache"] = past_key_values + mutable = ["cache"] + else: + mutable = False + + outputs = self.module.apply( + inputs, + jnp.array(input_ids, dtype="i4"), + jnp.array(attention_mask, dtype="i4"), + jnp.array(position_ids, dtype="i4"), + encoder_hidden_states, + encoder_attention_mask, + not train, + False, + output_attentions, + output_hidden_states, + return_dict, + rngs=rngs, + mutable=mutable, + ) + + # add updated cache to model output + if past_key_values is not None and return_dict: + outputs, past_key_values = outputs + outputs["past_key_values"] = unfreeze(past_key_values["cache"]) + return outputs + elif past_key_values is not None and not return_dict: + outputs, past_key_values = outputs + outputs = outputs[:1] + (unfreeze(past_key_values["cache"]),) + outputs[1:] + + return outputs + + +class LMCFlaxGPT2BlockCollection(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.blocks = [ + LMCFlaxGPT2Block(self.config, name=str(i), dtype=self.dtype) if i == self.config.att_layer_idx + else FlaxGPT2Block(self.config, name=str(i), dtype=self.dtype) + for i in range(self.config.num_hidden_layers) + ] + + def __call__( + self, + hidden_states, + attention_mask=None, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + deterministic: bool = True, + init_cache: bool = False, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ): + all_attentions = () if output_attentions else None + all_hidden_states = () if output_hidden_states else None + all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None + + for block in self.blocks: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = block( + hidden_states, + attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + deterministic=deterministic, + init_cache=init_cache, + output_attentions=output_attentions, + ) + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions += (layer_outputs[1],) + + if encoder_hidden_states is not None: + all_cross_attentions += (layer_outputs[2],) + + # this contains possible `None` values - `FlaxGPT2Module` will filter them out + outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions) + + return outputs + + +class LMCFlaxGPT2Module(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.embed_dim = self.config.hidden_size + + self.wte = nn.Embed( + self.config.vocab_size, + self.embed_dim, + embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range), + dtype=self.dtype, + ) + self.wpe = nn.Embed( + self.config.max_position_embeddings, + self.embed_dim, + embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range), + dtype=self.dtype, + ) + self.dropout = nn.Dropout(rate=self.config.embd_pdrop) + self.h = LMCFlaxGPT2BlockCollection(self.config, dtype=self.dtype) + self.ln_f = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype) + + def __call__( + self, + input_ids, + attention_mask, + position_ids, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + deterministic=True, + init_cache: bool = False, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ): + input_embeds = self.wte(input_ids.astype("i4")) + position_embeds = self.wpe(position_ids.astype("i4")) + + hidden_states = input_embeds + position_embeds + hidden_states = self.dropout(hidden_states, deterministic=deterministic) + + outputs = self.h( + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + deterministic=deterministic, + init_cache=init_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + hidden_states = self.ln_f(hidden_states) + + if output_hidden_states: + all_hidden_states = outputs[1] + (hidden_states,) + outputs = (hidden_states, all_hidden_states) + outputs[2:] + else: + outputs = (hidden_states,) + outputs[1:] + + if not return_dict: + return tuple(v for v in outputs if v is not None) + + return FlaxBaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + hidden_states=outputs[1], + attentions=outputs[2], + cross_attentions=outputs[3], + ) +class LMCFlaxGPT2Model(LMCFlaxGPT2PreTrainedModel): + module_class = LMCFlaxGPT2Module +class LMCFlaxGPT2LMHeadModule(nn.Module): + config: GPT2Config + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.transformer = LMCFlaxGPT2Module(self.config, dtype=self.dtype) + self.lm_head = nn.Dense( + self.config.vocab_size, + use_bias=False, + dtype=self.dtype, + kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range), + ) + + def __call__( + self, + input_ids, + attention_mask, + position_ids, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + deterministic: bool = True, + init_cache: bool = False, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ): + outputs = self.transformer( + input_ids, + attention_mask, + position_ids, + encoder_hidden_states, + encoder_attention_mask, + deterministic=deterministic, + init_cache=init_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + + if self.config.tie_word_embeddings: + shared_kernel = self.transformer.variables["params"]["wte"]["embedding"].T + lm_logits = self.lm_head.apply({"params": {"kernel": shared_kernel}}, hidden_states) + else: + lm_logits = self.lm_head(hidden_states) + + if not return_dict: + return (lm_logits,) + outputs[1:] + + return FlaxCausalLMOutputWithCrossAttentions( + logits=lm_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + cross_attentions=outputs.cross_attentions, + ) + + + +class LMCFlaxGPT2LMHeadModel(LMCFlaxGPT2PreTrainedModel): + module_class = LMCFlaxGPT2LMHeadModule + + def prepare_inputs_for_generation(self, input_ids, max_length, attention_mask: Optional[jax.Array] = None): + # initializing the cache + batch_size, seq_length = input_ids.shape + + past_key_values = self.init_cache(batch_size, max_length) + # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length. + # But since GPT2 uses a causal mask, those positions are masked anyways. + # Thus we can create a single static attention_mask here, which is more efficient for compilation + extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4") + if attention_mask is not None: + position_ids = attention_mask.cumsum(axis=-1) - 1 + extended_attention_mask = lax.dynamic_update_slice( + extended_attention_mask, attention_mask.astype("i4"), (0, 0) + ) + else: + position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length)) + + return { + "past_key_values": past_key_values, + "attention_mask": extended_attention_mask, + "position_ids": position_ids, + } + + def update_inputs_for_generation(self, model_outputs, model_kwargs): + model_kwargs["past_key_values"] = model_outputs.past_key_values + model_kwargs["position_ids"] = model_kwargs["position_ids"][:, -1:] + 1 + return model_kwargs diff --git a/src/lgmodeling/plot.py b/src/lgmodeling/plot.py new file mode 100644 index 0000000000000000000000000000000000000000..251a4f8466664792be4dd3b8fa730cefd9855b26 --- /dev/null +++ b/src/lgmodeling/plot.py @@ -0,0 +1,87 @@ +import argparse +import os +import json +import numpy as np +import matplotlib.pyplot as plt + +def calculate_loss_barrier(loss_curve): + return np.max(loss_curve) - (loss_curve[0] + loss_curve[-1]) / 2 + +def main(): + parser = argparse.ArgumentParser(description="Plot test loss and perplexity for best weight matching") + parser.add_argument("--file-1", type=str, required=True) + parser.add_argument("--file-2", type=str, required=True) + parser.add_argument("--file-3", type=str, required=True) + parser.add_argument("--output-dir", type=str, default=".", help="Directory to save output plot") + args = parser.parse_args() + + file_paths = [args.file_1, args.file_2, args.file_3] + xtick_labels = [ + ["Model 1", r"$\lambda$", "Model 2"], + ["Model 1", r"$\lambda$", "Model 3"], + ["Model 2", r"$\lambda$", "Model 3"], + ] + + data = [json.load(open(path, 'r')) for path in file_paths] + + num_points = len(data[0]["Naive"]["Test Loss"],) + lambda_values = np.linspace(0, 1, num_points) + + plt.rcParams.update({ + "font.family": "serif", + 'legend.frameon': False, + 'lines.linewidth': 2, + }) + + colors = ["steelblue", "lightsalmon"] + plt.style.use('tableau-colorblind10') + fig, axs = plt.subplots(2, 3, figsize=(13, 6)) + + FONT_SMALL = 11 + FONT_MEDIUM = 13 + FONT_LARGE = 16 + + # Determine best clever index (lowest test loss barrier) for each file + # best_indices = [] + # for d in data: + # best_idx = 0 + # min_barrier = float("inf") + # for idx, loss_curve in enumerate(d["test_loss_interp_clever_list"]): + # barrier = calculate_loss_barrier(loss_curve) + # if barrier < min_barrier: + # min_barrier = barrier + # best_idx = idx + # best_indices.append(best_idx) + + # Plot each model pair + for i in range(3): + d = data[i] + # best_idx = best_indices[i] + # Row 0: Test Loss + ax = axs[0, i] + ax.plot(lambda_values, d["Naive"]["Test Loss"], label="Naive", color=colors[0]) + ax.plot(lambda_values, d["permu_head_init_ortho_opt"]["Test Loss"], label="Match", color=colors[1]) + ax.set_xticks([0, 0.5, 1]) + ax.set_xticklabels(xtick_labels[i], fontsize=FONT_MEDIUM) + ax.tick_params(axis='y', labelsize=FONT_MEDIUM) + if i == 0: + ax.set_ylabel("Test Loss", fontsize=FONT_LARGE, labelpad=15) + ax.legend(loc='upper left', fontsize=FONT_SMALL) + + # Row 1: Test Perplexity + ax = axs[1, i] + ax.plot(lambda_values, d["Naive"]["Test PPL"], label="Naive", color=colors[0]) + ax.plot(lambda_values, d["permu_head_init_ortho_opt"]["Test PPL"], label="Match", color=colors[1]) + ax.set_xticks([0, 0.5, 1]) + ax.set_xticklabels(xtick_labels[i], fontsize=FONT_MEDIUM) + ax.tick_params(axis='y', labelsize=FONT_MEDIUM) + if i == 0: + ax.set_ylabel("Test BPC", fontsize=FONT_LARGE, labelpad=15) + + plt.tight_layout(w_pad=4.0, h_pad=2.5) + output_path = args.output_dir + plt.savefig(output_path) + print(f"Saved plot to {output_path}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/lgmodeling/train_model.py b/src/lgmodeling/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..2c086bd8bf67c9404789e6fb55f192570211778f --- /dev/null +++ b/src/lgmodeling/train_model.py @@ -0,0 +1,294 @@ +import argparse +import jax +import os +import optax +import torch +import wandb +import math +import time +import itertools +import numpy as np +import jax.numpy as jnp +from datetime import timedelta +from typing import Any, Dict, List +from datasets import Dataset +from tqdm import tqdm +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.training.common_utils import get_metrics, onehot, shard +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.gpt2.modeling_flax_gpt2 import GPT2Config +from lmc_model import LMCFlaxGPT2LMHeadModel +from data_utils import get_lm_corpus +from jax import debug +os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" + +# fmt: off + +# fmt: on + + +def prepare_lm_batch(data: torch.Tensor, target: torch.Tensor) -> Dict[str, Any]: + """ + Convert and shard a language modeling batch from PyTorch to JAX. + Args: + data (torch.Tensor): Input data of shape (seq_len, batch) + target (torch.Tensor): Target data of shape (seq_len, batch) + Returns: + Dict[str, jnp.ndarray]: Dict with 'data' and 'target', both sharded + with shape (n_devices, batch_per_device, seq_len) + """ + # Transpose to (batch, seq_len), then convert to jnp arrays + input_ids = jnp.array(data.T) + target = jnp.array(target.T) + # Shard across devices + return {'input_ids': shard(input_ids),'target': shard(target)} + + +def decay_mask_fn(params): + flat_params = flatten_dict(params) + flat_mask = { + path: (path[-1] != "bias" and path[-2:] not in [("ln_1", "scale"), ("ln_2", "scale"), ("ln_f", "scale")]) + for path in flat_params + } + return unflatten_dict(flat_mask) + + +def main(args: argparse.Namespace): + os.makedirs(args.wandb_run_dir, exist_ok=True) + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + name=f"lr{args.learning_rate}-{args.position_embeddings}-step{args.max_step}-warm{args.warmup_step}" + f"-size{args.batch_size}-layer{args.n_layer}-embd{args.n_embd}-heads{args.n_head}" + f"-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", + save_code=True, + ) + save_path = os.path.join(args.model_save_dir,wandb.run.name) + wandb.config = dict(vars(args)) + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + corpus = get_lm_corpus(args.data_path, args.dataset) + ntokens = len(corpus.vocab) + args.n_token = ntokens + eval_batch_size = 12 + tr_iter = corpus.get_iterator('train', args.batch_size, args.tgt_len, ext_len=args.ext_len) + va_iter = corpus.get_iterator('valid', eval_batch_size, args.eval_tgt_len, ext_len=args.ext_len) + te_iter = corpus.get_iterator('test', eval_batch_size, args.eval_tgt_len, ext_len=args.ext_len) + model_config = GPT2Config() + model_config.position_embeddings = args.position_embeddings + model_config.rotary_dim = args.rotary_dim + model_config.vocab_size = args.n_token + model_config.n_positions = args.tgt_len + model_config.n_ctx = args.tgt_len + model_config.bos_token_id = args.n_token + model_config.eos_token_id = args.n_token + model_config.n_layer = args.n_layer + model_config.n_head = args.n_head + model_config.n_embd = args.n_embd + model_config.n_inner = args.n_inner + model_config.lmc_layer_indices = args.lmc_layer_indices + model_config.num_routed_experts = args.num_routed_experts + model_config.num_shared_experts = args.num_shared_experts + model_config.topk = args.topk + model_config.routed_scaling_factor = args.routed_scaling_factor + model_config.q_lora_rank = args.q_lora_rank + model_config.qk_nope_head_dim = args.qk_nope_head_dim + model_config.qk_rope_head_dim = args.qk_rope_head_dim + model_config.kv_lora_rank = args.kv_lora_rank + model_config.v_head_dim = args.v_head_dim + model_config.rope_scaling = None + model_config.attention_bias = args.attention_bias + model = LMCFlaxGPT2LMHeadModel(model_config,input_shape=(1, args.tgt_len),seed=0,dtype=jnp.dtype(args.dtype),) + model.config.save_pretrained(save_path) + num_train_steps = args.max_step + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=0.0, + peak_value=args.learning_rate, + warmup_steps=args.warmup_step, + decay_steps=args.max_step, + end_value=args.eta_min, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay_rate, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip train step to {state.step}") + latest_train_step = state.step + def train_step(state, batch, dropout_rng): + dropout_rng, new_dropout_rng = jax.random.split(dropout_rng) + def loss_fn(params): + labels = batch.pop("target") + logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0] + loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])).mean() + return loss + grad_fn = jax.value_and_grad(loss_fn) + loss, grads = grad_fn(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + new_state = state.apply_gradients(grads=grads) + metrics = {"loss": loss,"learning_rate": lr_schedule(state.step)} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return new_state, metrics, new_dropout_rng + def eval_step(state, batch): + labels = batch.pop("target") + logits = model(**batch, params=state.params, train=False)[0] + loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])).mean() + metrics = {"eval_loss": loss} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + state = replicate(state) + train_metrics_stack = [] + train_step = int(jax.device_get(unreplicate(state.step))) + train_loss = 0.0 + best_val_loss = float("inf") + log_start_time = time.time() + eval_start_time = time.time() + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in itertools.count(start=1): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_iter = tr_iter.get_varlen_iter() if getattr(args, "varlen", False) else tr_iter + train_metrics_stack = [] + for batch_idx, (data, target, seq_len) in enumerate(tqdm(train_iter)): + if train_step >= args.max_step: + break + # Prepare and shard batch + batch = prepare_lm_batch(data, target) + # Run train step + state, train_metric, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics_stack.append(train_metric) + train_step += 1 + # Logging + if train_step % args.logging_frequency == 0: + train_metrics = get_metrics(train_metrics_stack) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(lambda x: x.mean(), train_metrics) + train_metrics_stack = [] + loss = float(train_metrics["loss"]) + ppl = math.exp(loss) + bpc = loss/math.log(2) + curr_lr = float(lr_schedule(train_step)) + elapsed = time.time() - log_start_time + if(args.dataset in ["wt103","lm1b"]): + print( + f"| epoch {epoch:3d} step {train_step:8d} | " + f"{batch_idx+1:6d} batches | lr {curr_lr:.3g} " + f"| ms/batch {elapsed * 1000 / args.logging_frequency:5.2f} | " + f"loss {loss:5.2f} | ppl {ppl:9.3f}" + ) + wandb.log({"loss": loss,"ppl": ppl,"learning_rate": curr_lr}, step=train_step) + elif(args.dataset in ["enwik8","text8"]): + print( + f"| epoch {epoch:3d} step {train_step:8d} | " + f"{batch_idx+1:6d} batches | lr {curr_lr:.3g} " + f"| ms/batch {elapsed * 1000 / args.logging_frequency:5.2f} | " + f"loss {loss:5.2f} | bpc {bpc:9.3f}" + ) + wandb.log({"loss": loss,"bpc": bpc,"learning_rate": curr_lr}, step=train_step) + log_start_time = time.time() + # Evaluation + if train_step % args.eval_frequency == 0: + eval_results = [] + for eval_data, eval_target, _ in va_iter: + eval_batch = prepare_lm_batch(eval_data, eval_target) + eval_metric = parallel_eval_step(state, eval_batch) + eval_results.append(eval_metric) + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(lambda x: x.mean(), eval_metrics) + + val_loss = float(eval_metrics["eval_loss"]) + val_ppl = math.exp(val_loss) + val_bpc = val_loss/math.log(2) + print("-" * 100) + if(args.dataset in ["wt103","lm1b"]): + print( + f"| Eval {train_step // args.eval_frequency:3d} at step {train_step:8d} | " + f"time: {time.time() - eval_start_time:5.2f}s | " + f"valid loss {val_loss:5.2f} | valid ppl {val_ppl:9.3f}" + ) + wandb.log({"eval_loss": val_loss,"eval_ppl": val_ppl}, step=train_step) + elif(args.dataset in ["enwik8","text8"]): + print( + f"| Eval {train_step // args.eval_frequency:3d} at step {train_step:8d} | " + f"time: {time.time() - eval_start_time:5.2f}s | " + f"valid loss {val_loss:5.2f} | valid bpc {val_bpc:9.3f}" + ) + wandb.log({"eval_loss": val_loss,"eval_bpc": val_bpc}, step=train_step) + print("-" * 100) + # Save best checkpoint + if val_loss < best_val_loss: + best_val_loss = val_loss + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=train_step,prefix="best_",keep=1) + print(f"Best model saved at step {train_step}") + eval_start_time = time.time() + # Periodic checkpoint + if train_step % args.save_frequency == 0: + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=train_step,prefix="last_",keep=1) + print(f"Checkpoint saved at step {save_path}") + if train_step >= args.max_step: + print("-" * 100) + print("End of training") + break + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--rotary-dim", type=int,default=0,help=' Number of dimensions in the embedding that Rotary Position Embedding is applied to..') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--n_layer', type=int, default=12,help='number of total layers') + parser.add_argument('--n_head', type=int, default=10,help='number of heads') + parser.add_argument('--n_embd', type=int, default=500,help='model dimension') + parser.add_argument('--n_inner', type=int, default=1000,help='inner dimension in FF') + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--seed", type=int, default=0, help="random seed for RNG state") + parser.add_argument("--data-path", type=str, default="", help="train datset paths (multiple paths)") + parser.add_argument('--dataset', type=str, default='wt103',choices=['wt103', 'lm1b', 'enwik8', 'text8'],help='dataset name') + parser.add_argument('--max_step', type=int, default=500000,help='upper epoch limit') + parser.add_argument('--warmup_step', type=int, default=2000,help='upper epoch limit') + parser.add_argument("--batch-size", type=int, default=96, help="train, eval batch size (batch size will be devided by device count)") + parser.add_argument('--tgt_len', type=int, default=256,help='number of tokens to predict') + parser.add_argument('--eval_tgt_len', type=int, default=256,help='number of tokens to predict for evaluation') + parser.add_argument('--ext_len', type=int, default=0,help='length of the extended context') + parser.add_argument('--mem_len', type=int, default=0,help='length of the retained previous heads') + parser.add_argument("--learning-rate", type=float, default=0.00025, help="learning rate") + parser.add_argument("--weight-decay-rate", type=float, default=0.01, help="weight deacy rate for lr scheduler") + parser.add_argument('--eta_min', type=float, default=1.0e-8,help='min learning rate for cosine scheduler') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default="", help="wandb entity for logging") + parser.add_argument("--wandb-group", default="", help="wandb group for logging") + parser.add_argument("--wandb-project", default="GPT2-Wikitext103", help="wandb project name for logging") + parser.add_argument("--wandb-run-dir", default=".wandb", help="wandb run dir") + parser.add_argument("--logging-frequency", type=int, default=200, help="do logging every logging_frequency step") + parser.add_argument("--eval-frequency", type=int, default=4000, help="do evalution every eval_frequency step") + parser.add_argument("--save-frequency", type=int, default=4000, help="do saving checkpoint every save_frequencey step") + parser.add_argument("--model-save-dir", type=str, default="artifacts/", help="checkpoint saving dir") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + main(parser.parse_args()) \ No newline at end of file diff --git a/src/lgmodeling/utils.py b/src/lgmodeling/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0f425c2e9fad683b43d981a2458ed6aa5973dbcd --- /dev/null +++ b/src/lgmodeling/utils.py @@ -0,0 +1,118 @@ +import operator +import re +import time +from contextlib import contextmanager + +import jax.numpy as jnp +from flax import traverse_util +from flax.core import freeze, unfreeze +from jax import random +from jax.tree_util import tree_reduce, tree_map + +rngmix = lambda rng, x: random.fold_in(rng, hash(x) % 2**16) + +@contextmanager +def timeblock(name): + start = time.time() + try: + yield + finally: + end = time.time() + print(f"{name} took {end - start:.5f} seconds") + +class RngPooper: + """A stateful wrapper around stateless random.PRNGKey's.""" + + def __init__(self, init_rng): + self.rng = init_rng + + def poop(self): + self.rng, rng_key = random.split(self.rng) + return rng_key + +def l1prox(x, alpha): + return jnp.sign(x) * jnp.maximum(0, jnp.abs(x) - alpha) + +def ec2_get_instance_type(): + # See also https://stackoverflow.com/questions/51486405/aws-ec2-command-line-display-instance-type/51486782 + return open("/sys/devices/virtual/dmi/id/product_name").read().strip() + +# Utilities for dealing with flax model parameters +def partition(pred, iterable): + trues = [] + falses = [] + for item in iterable: + if pred(item): + trues.append(item) + else: + falses.append(item) + return trues, falses + +def partition_dict(pred, d): + trues = {} + falses = {} + for k, v in d.items(): + if pred(k): + trues[k] = v + else: + falses[k] = v + return trues, falses + +def flatten_params(params): + return {"/".join(k): v for k, v in traverse_util.flatten_dict(unfreeze(params)).items()} + +def unflatten_params(flat_params): + return freeze( + traverse_util.unflatten_dict({tuple(k.split("/")): v + for k, v in flat_params.items()})) + +def merge_params(a, b): + return unflatten_params({**a, **b}) + +def kmatch(pattern, key): + regex = "^" + i = 0 + while i < len(pattern): + if pattern[i] == "*": + if i + 1 < len(pattern) and pattern[i + 1] == "*": + regex += "(.*)" + i += 2 + else: + regex += "([^\/]*)" + i += 1 + else: + regex += pattern[i] + i += 1 + regex += "$" + return re.fullmatch(regex, key) + +assert kmatch("*", "a") is not None +assert kmatch("*", "a").group(0) == "a" +assert kmatch("*", "a").group(1) == "a" +assert kmatch("abc", "def") is None +assert kmatch("abc/*/ghi", "abc/def/ghi").group(1) == "def" +assert kmatch("abc/**/jkl", "abc/def/ghi/jkl").group(1) == "def/ghi" +assert kmatch("abc/*/jkl", "abc/def/ghi/jkl") is None +assert kmatch("**/*", "abc/def/ghi/jkl").group(1) == "abc/def/ghi" +assert kmatch("**/*", "abc/def/ghi/jkl").group(2) == "jkl" + +def lerp(lam, t1, t2): + return tree_map(lambda a, b: (1 - lam) * a + lam * b, t1, t2) + +def tree_norm(t): + return jnp.sqrt(tree_reduce(operator.add, tree_map(lambda x: jnp.sum(x**2), t))) + +def tree_l2(t1, t2): + return tree_norm(tree_map(lambda x, y: x - y, t1, t2)) + +def slerp(lam, t1, t2): + # See https://en.wikipedia.org/wiki/Slerp + om = jnp.arccos( + tree_reduce(operator.add, tree_map(lambda x, y: jnp.sum(x * y), t1, t2)) / + (tree_norm(t1) * tree_norm(t2))) + sinom = jnp.sin(om) + return tree_map( + lambda x, y: jnp.sin((1 - lam) * om) / sinom * x + jnp.sin(lam * om) / sinom * y, + t1, + t2, + ) diff --git a/src/lgmodeling/vocabulary.py b/src/lgmodeling/vocabulary.py new file mode 100644 index 0000000000000000000000000000000000000000..b6b8249b1f5ef1362183f45ace645b41162fe2f4 --- /dev/null +++ b/src/lgmodeling/vocabulary.py @@ -0,0 +1,163 @@ +import os +from collections import Counter, OrderedDict + +import torch + +class Vocab(object): + def __init__(self, special=[], min_freq=0, max_size=None, lower_case=True, + delimiter=None, vocab_file=None): + self.counter = Counter() + self.special = special + self.min_freq = min_freq + self.max_size = max_size + self.lower_case = lower_case + self.delimiter = delimiter + self.vocab_file = vocab_file + + def tokenize(self, line, add_eos=False, add_double_eos=False): + line = line.strip() + # convert to lower case + if self.lower_case: + line = line.lower() + + # empty delimiter '' will evaluate False + if self.delimiter == '': + symbols = line + else: + symbols = line.split(self.delimiter) + + if add_double_eos: # lm1b + return [''] + symbols + [''] + elif add_eos: + return symbols + [''] + else: + return symbols + + def count_file(self, path, verbose=False, add_eos=False): + if verbose: print('counting file {} ...'.format(path)) + assert os.path.exists(path) + + sents = [] + with open(path, 'r', encoding='utf-8') as f: + for idx, line in enumerate(f): + if verbose and idx > 0 and idx % 500000 == 0: + print(' line {}'.format(idx)) + symbols = self.tokenize(line, add_eos=add_eos) + self.counter.update(symbols) + sents.append(symbols) + + return sents + + def count_sents(self, sents, verbose=False): + """ + sents : a list of sentences, each a list of tokenized symbols + """ + if verbose: print('counting {} sents ...'.format(len(sents))) + for idx, symbols in enumerate(sents): + if verbose and idx > 0 and idx % 500000 == 0: + print(' line {}'.format(idx)) + self.counter.update(symbols) + + def _build_from_file(self, vocab_file): + self.idx2sym = [] + self.sym2idx = OrderedDict() + + with open(vocab_file, 'r', encoding='utf-8') as f: + for line in f: + symb = line.strip().split()[0] + self.add_symbol(symb) + self.unk_idx = self.sym2idx[''] + + def build_vocab(self): + if self.vocab_file: + print('building vocab from {}'.format(self.vocab_file)) + self._build_from_file(self.vocab_file) + print('final vocab size {}'.format(len(self))) + else: + print('building vocab with min_freq={}, max_size={}'.format( + self.min_freq, self.max_size)) + self.idx2sym = [] + self.sym2idx = OrderedDict() + + for sym in self.special: + self.add_special(sym) + + for sym, cnt in self.counter.most_common(self.max_size): + if cnt < self.min_freq: break + self.add_symbol(sym) + + print('final vocab size {} from {} unique tokens'.format( + len(self), len(self.counter))) + + def encode_file(self, path, ordered=False, verbose=False, add_eos=True, + add_double_eos=False): + if verbose: print('encoding file {} ...'.format(path)) + assert os.path.exists(path) + encoded = [] + with open(path, 'r', encoding='utf-8') as f: + for idx, line in enumerate(f): + if verbose and idx > 0 and idx % 500000 == 0: + print(' line {}'.format(idx)) + symbols = self.tokenize(line, add_eos=add_eos, + add_double_eos=add_double_eos) + encoded.append(self.convert_to_tensor(symbols)) + + if ordered: + encoded = torch.cat(encoded) + + return encoded + + def encode_sents(self, sents, ordered=False, verbose=False): + if verbose: print('encoding {} sents ...'.format(len(sents))) + encoded = [] + for idx, symbols in enumerate(sents): + if verbose and idx > 0 and idx % 500000 == 0: + print(' line {}'.format(idx)) + encoded.append(self.convert_to_tensor(symbols)) + + if ordered: + encoded = torch.cat(encoded) + + return encoded + + def add_special(self, sym): + if sym not in self.sym2idx: + self.idx2sym.append(sym) + self.sym2idx[sym] = len(self.idx2sym) - 1 + setattr(self, '{}_idx'.format(sym.strip('<>')), self.sym2idx[sym]) + + def add_symbol(self, sym): + if sym not in self.sym2idx: + self.idx2sym.append(sym) + self.sym2idx[sym] = len(self.idx2sym) - 1 + + def get_sym(self, idx): + assert 0 <= idx < len(self), 'Index {} out of range'.format(idx) + return self.idx2sym[idx] + + def get_idx(self, sym): + if sym in self.sym2idx: + return self.sym2idx[sym] + else: + # print('encounter unk {}'.format(sym)) + assert '' not in sym + assert hasattr(self, 'unk_idx') + return self.sym2idx.get(sym, self.unk_idx) + + def get_symbols(self, indices): + return [self.get_sym(idx) for idx in indices] + + def get_indices(self, symbols): + return [self.get_idx(sym) for sym in symbols] + + def convert_to_tensor(self, symbols): + return torch.LongTensor(self.get_indices(symbols)) + + def convert_to_sent(self, indices, exclude=None): + if exclude is None: + return ' '.join([self.get_sym(idx) for idx in indices]) + else: + return ' '.join([self.get_sym(idx) for idx in indices if idx not in exclude]) + + def __len__(self): + return len(self.idx2sym) diff --git a/src/matching_utils.py b/src/matching_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3ad0e01ee46465e0bba42bdfe8e39124082d3add --- /dev/null +++ b/src/matching_utils.py @@ -0,0 +1,1347 @@ +from collections import defaultdict +from typing import NamedTuple +from flax.core import freeze, unfreeze +import jax.numpy as jnp +from jax import random, tree_util, jit, grad, value_and_grad +from scipy.optimize import linear_sum_assignment, minimize +import numpy as np +import matplotlib.pyplot as plt +import time +import os +import copy +import jax +import jax.lax as lax +import jax.nn as nn +import jax + +'''num_heads = 4 + #print all layers + layer_paths = [] + def collect_layer_paths(path, value): + # Convert path to a readable string by joining path keys + path_str = '/'.join([str(p.key) for p in path]) + shape = value.shape + layer_paths.append((path_str, shape)) + + jax.tree_util.tree_map_with_path(collect_layer_paths, pretrained_params) + print("Layers of the model:") + for path, shape in layer_paths: + print(f" {path} {shape}")''' +'''example output + Layers of the model: + Conv_0/bias (32,) + Conv_0/kernel (4, 4, 3, 32) + Dense_0/bias (10,) + Dense_0/kernel (32, 10) + TransformerEncoderLayer_0/Dense_0/bias (128,) + TransformerEncoderLayer_0/Dense_0/kernel (32, 128) + TransformerEncoderLayer_0/Dense_1/bias (32,) + TransformerEncoderLayer_0/Dense_1/kernel (128, 32) + TransformerEncoderLayer_0/LayerNorm_0/bias (32,) + TransformerEncoderLayer_0/LayerNorm_0/scale (32,) + TransformerEncoderLayer_0/LayerNorm_1/bias (32,) + TransformerEncoderLayer_0/LayerNorm_1/scale (32,) + TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/key/bias (4, 8) + TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/key/kernel (32, 4, 8) #Key projections for 4 attention heads, each with a dimension of 8 (4 heads x 8 = 32, matching the model's hidden size). + TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/out/bias (32,) + TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/out/kernel (4, 8, 32) + TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/query/bias (4, 8) + TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/query/kernel (32, 4, 8) + TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/value/bias (4, 8) + TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/value/kernel (32, 4, 8) + TransformerEncoderLayer_1/Dense_0/bias (128,) + TransformerEncoderLayer_1/Dense_0/kernel (32, 128) + TransformerEncoderLayer_1/Dense_1/bias (32,) + TransformerEncoderLayer_1/Dense_1/kernel (128, 32) + TransformerEncoderLayer_1/LayerNorm_0/bias (32,) + TransformerEncoderLayer_1/LayerNorm_0/scale (32,) + TransformerEncoderLayer_1/LayerNorm_1/bias (32,) + TransformerEncoderLayer_1/LayerNorm_1/scale (32,) + TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/key/bias (4, 8) + TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/key/kernel (32, 4, 8) + TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/out/bias (32,) + TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/out/kernel (4, 8, 32) + TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/query/bias (4, 8) + TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/query/kernel (32, 4, 8) + TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/value/bias (4, 8) + TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/value/kernel (32, 4, 8) + TransformerEncoderLayer_2/Dense_0/bias (128,) + TransformerEncoderLayer_2/Dense_0/kernel (32, 128) + TransformerEncoderLayer_2/Dense_1/bias (32,) + TransformerEncoderLayer_2/Dense_1/kernel (128, 32) + TransformerEncoderLayer_2/LayerNorm_0/bias (32,) + TransformerEncoderLayer_2/LayerNorm_0/scale (32,) + TransformerEncoderLayer_2/LayerNorm_1/bias (32,) + TransformerEncoderLayer_2/LayerNorm_1/scale (32,) + TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/key/bias (4, 8) + TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/key/kernel (32, 4, 8) + TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/out/bias (32,) + TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/out/kernel (4, 8, 32) + TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/query/bias (4, 8) + TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/query/kernel (32, 4, 8) + TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/value/bias (4, 8) + TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/value/kernel (32, 4, 8) + cls_token (1, 1, 32) + pos_embedding (1, 65, 32) +''' + +def to_numpy(x): + if isinstance(x, jnp.ndarray): + return np.array(x) + return np.array(x) + +@jit +def compute_objective_jax(A, X, X_prime, Y, Y_prime, cond_threshold=1e6): + cond = jnp.linalg.cond(A) + def safe_obj(): + A_inv = jnp.linalg.inv(A) + term1 = X - X_prime @ A.T + term2 = Y - Y_prime @ A_inv + return jnp.sum(term1**2) + jnp.sum(term2**2) + return lax.cond(cond > cond_threshold, lambda: jnp.inf, safe_obj) + +compute_value_and_grad_jax = jit(value_and_grad(compute_objective_jax)) + +def solve_orthogonal(X, X_prime, Y, Y_prime): + B = X.T @ X_prime + Y.T @ Y_prime + U, _, Vt = np.linalg.svd(B) + return U @ Vt + +def optimize_alignment(A_init, X, X_prime, Y, Y_prime, max_iter=5000): + objective_values = [] + grad_norms = [] + condition_nums = [] + + def obj_fn(flat_A): + A = flat_A.reshape(A_init.shape) + obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime)) + return float(obj), np.array(grad_val).flatten() + + def callback(flat_A): + A = flat_A.reshape(A_init.shape) + obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime)) + grad_norm = jnp.linalg.norm(grad_val, 'fro') + cond = jnp.linalg.cond(jnp.array(A)) + objective_values.append(float(obj)) + grad_norms.append(float(grad_norm)) + condition_nums.append(float(cond)) + + res = minimize(obj_fn, A_init.flatten(), jac=True, method='L-BFGS-B', options={'maxiter': max_iter}, callback=callback) + A_opt = res.x.reshape(A_init.shape) + return A_opt, objective_values, grad_norms, condition_nums + +def get_nested_item(d, keys): + """Accesses a nested dictionary item using a tuple of keys.""" + for key in keys: + d = d[key] + return d + +# Helper functions for parameter extraction and reshaping +def extract_attention_params(params, layer_idx): + """ + Extracts MHA parameters from different, known model structures in a compatible way. + + This function detects the model type and constructs the correct path to the + attention parameters for a given layer index. + + Args: + params: The Flax parameter tree. + layer_idx: The integer index of the transformer layer. + + Returns: + A tuple containing: + - A flat tuple of the MHA tensors: (key, key_bias, query, query_bias, value, value_bias, out, out_bias). + - A tuple representing the nested path to the MHA block, for use in updates. + """ + # Detect vit-jax style: params['Transformer']['encoderblock_...'] + if 'Transformer' in params and f'encoderblock_{layer_idx}' in params['Transformer']: + mha_path = ('Transformer', f'encoderblock_{layer_idx}', 'MultiHeadDotProductAttention_0') + # Detect cifar_vit style: params['TransformerEncoderLayer_...'] + elif f'TransformerEncoderLayer_{layer_idx}' in params: + mha_path = (f'TransformerEncoderLayer_{layer_idx}', 'MultiHeadDotProductAttention_0') + else: + raise KeyError(f"Could not find a known path for attention layer {layer_idx} in the provided params.") + + attention_block = get_nested_item(params, mha_path) + + key_k, key_b = attention_block['key']['kernel'], attention_block['key']['bias'] + query_k, query_b = attention_block['query']['kernel'], attention_block['query']['bias'] + value_k, value_b = attention_block['value']['kernel'], attention_block['value']['bias'] + out_k, out_b = attention_block['out']['kernel'], attention_block['out']['bias'] + + return (key_k, key_b, query_k, query_b, value_k, value_b, out_k, out_b), mha_path + +def reshape_to_per_head(params, num_heads): + """ + Reshapes batched attention parameters into a list of per-head parameters. + + This function assumes a specific shape convention for the input weight and + bias tensors, which is common in Flax/Linen implementations. + + Args: + params (dict): A dictionary containing the attention parameters. + Expected keys and tensor shapes are: + - 'query': Weight tensor of shape (D, num_heads, d_k) + - 'query_bias': Bias tensor of shape (num_heads, d_k) + - 'key': Weight tensor of shape (D, num_heads, d_k) + - 'key_bias': Bias tensor of shape (num_heads, d_k) + - 'value': Weight tensor of shape (D, num_heads, d_v) + - 'value_bias': Bias tensor of shape (num_heads, d_v) + - 'out': Weight tensor of shape (num_heads, d_v, D) + num_heads (int): The number of attention heads. + + Returns: + A tuple containing lists of per-head parameters: + (W_Q, b_Q, W_K, b_K, W_V, b_V, W_O) + """ + + query_kernel = params['query'] + assert query_kernel.ndim == 3, f"Expected query weights to be 3D, but got shape {query_kernel.shape}" + assert query_kernel.shape[1] == num_heads, ( + f"The second dimension of the query weight tensor should be num_heads ({num_heads}), " + f"but got shape {query_kernel.shape}. Please verify your model's parameter shape convention." + ) + + W_Q = [params['query'][:, i, :] for i in range(num_heads)] + b_Q = [params['query_bias'][i, :] for i in range(num_heads)] + W_K = [params['key'][:, i, :] for i in range(num_heads)] + b_K = [params['key_bias'][i, :] for i in range(num_heads)] + W_V = [params['value'][:, i, :] for i in range(num_heads)] + b_V = [params['value_bias'][i, :] for i in range(num_heads)] + W_O = [params['out'][i, :, :] for i in range(num_heads)] + + return W_Q, b_Q, W_K, b_K, W_V, b_V, W_O + +def compute_extended_weights(W, b): + return jnp.vstack([jnp.array(W), jnp.array(b).reshape(1, -1)]) + +# Helper function to plot multiple curves +def plot_multiple_curves(data_list, title, xlabel, ylabel, labels, save_path): + plt.figure() + for data, label in zip(data_list, labels): + label = f"{label} ({data[-1]:.4f})" + plt.plot(data, label=label) + plt.title(title) + plt.xlabel(xlabel) + plt.ylabel(ylabel) + plt.legend(loc='center left', bbox_to_anchor=(1, 0.5)) + plt.savefig(save_path, bbox_inches='tight') + plt.close() + +# Stage 1 Function: Find Heads Permutation (Data-Dependent) +# Version 1: Using post-softmax probabilities +def compute_cost_matrix_postsoftmax(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, activations_a, activations_b, alpha=0.5, epsilon=1e-8): + """ + Computes the cost matrix using post-softmax probabilities and model-specific activations. + """ + B_a, L_a, D_a = activations_a.shape + B_b, L_b, D_b = activations_b.shape + assert B_a == B_b and L_a == L_b and D_a == D_b, "Activations for both models must have the same shape." + + # Augment activations for each model separately + X_tilde_a = jnp.concatenate([activations_a, jnp.ones((B_a, L_a, 1))], axis=-1) + X_tilde_b = jnp.concatenate([activations_b, jnp.ones((B_b, L_b, 1))], axis=-1) + + d_head = W_Q_a[0].shape[1] + sqrt_d = jnp.sqrt(float(d_head)) + C = np.zeros((num_heads, num_heads)) + + # Pre-compute flattened outputs for model A + P_flat_a, V_flat_a = [], [] + for i in range(num_heads): + tilde_W_Q_a_i = compute_extended_weights(W_Q_a[i], b_Q_a[i]) + tilde_W_K_a_i = compute_extended_weights(W_K_a[i], b_K_a[i]) + tilde_W_V_a_i = compute_extended_weights(W_V_a[i], b_V_a[i]) + Q_a_i = X_tilde_a @ tilde_W_Q_a_i + K_a_i = X_tilde_a @ tilde_W_K_a_i + S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d + P_a_i = nn.softmax(S_a_i, axis=-1) + P_flat_a.append(P_a_i.flatten()) + V_a_i = (X_tilde_a @ tilde_W_V_a_i) @ W_O_a[i] + V_flat_a.append(V_a_i.flatten()) + + # Pre-compute flattened outputs for model B + P_flat_b, V_flat_b = [], [] + for j in range(num_heads): + tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j]) + tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j]) + tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j]) + Q_b_j = X_tilde_b @ tilde_W_Q_b_j + K_b_j = X_tilde_b @ tilde_W_K_b_j + S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d + P_b_j = nn.softmax(S_b_j, axis=-1) + P_flat_b.append(P_b_j.flatten()) + V_b_j = (X_tilde_b @ tilde_W_V_b_j) @ W_O_b[j] + V_flat_b.append(V_b_j.flatten()) + + # Compute cost matrix from pre-computed values + for i in range(num_heads): + for j in range(num_heads): + # Cosine similarity for P (post-softmax probabilities) + dot_P = jnp.dot(P_flat_a[i], P_flat_b[j]) + norm_P_a = jnp.linalg.norm(P_flat_a[i]) + norm_P_b = jnp.linalg.norm(P_flat_b[j]) + cost_P = 1.0 - (dot_P / (norm_P_a * norm_P_b + epsilon)) + + # Cosine similarity for V (value-projections) + dot_V = jnp.dot(V_flat_a[i], V_flat_b[j]) + norm_V_a = jnp.linalg.norm(V_flat_a[i]) + norm_V_b = jnp.linalg.norm(V_flat_b[j]) + cost_V = 1.0 - (dot_V / (norm_V_a * norm_V_b + epsilon)) + + C[i, j] = alpha * cost_P + (1 - alpha) * cost_V + + return C + +# Version 2: Using pre-softmax scores +def compute_cost_matrix_presoftmax(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, activations_a, activations_b, alpha=0.5, epsilon=1e-8): + """ + Computes the cost matrix for attention head permutation using model-specific activations. + """ + B_a, L_a, D_a = activations_a.shape + B_b, L_b, D_b = activations_b.shape + assert B_a == B_b and L_a == L_b and D_a == D_b, "Activations for both models must have the same shape." + + # Augment activations for model A + ones_col_a = jnp.ones((B_a, L_a, 1)) + X_tilde_a = jnp.concatenate([activations_a, ones_col_a], axis=-1) + + # Augment activations for model B + ones_col_b = jnp.ones((B_b, L_b, 1)) + X_tilde_b = jnp.concatenate([activations_b, ones_col_b], axis=-1) + + d_head = W_Q_a[0].shape[1] + sqrt_d = jnp.sqrt(float(d_head)) + C = np.zeros((num_heads, num_heads)) + + # Pre-compute all head outputs for model A + S_bar_flat_a, V_flat_a = [], [] + for i in range(num_heads): + tilde_W_Q_a_i = compute_extended_weights(W_Q_a[i], b_Q_a[i]) + tilde_W_K_a_i = compute_extended_weights(W_K_a[i], b_K_a[i]) + tilde_W_V_a_i = compute_extended_weights(W_V_a[i], b_V_a[i]) + Q_a_i = X_tilde_a @ tilde_W_Q_a_i + K_a_i = X_tilde_a @ tilde_W_K_a_i + S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d + S_bar_a_i = S_a_i - jnp.mean(S_a_i, axis=2, keepdims=True) + S_bar_flat_a.append(S_bar_a_i.flatten()) + V_tilde_a_i = X_tilde_a @ tilde_W_V_a_i + V_a_i = V_tilde_a_i @ W_O_a[i] + V_flat_a.append(V_a_i.flatten()) + + # Pre-compute all head outputs for model B + S_bar_flat_b, V_flat_b = [], [] + for j in range(num_heads): + tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j]) + tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j]) + tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j]) + Q_b_j = X_tilde_b @ tilde_W_Q_b_j + K_b_j = X_tilde_b @ tilde_W_K_b_j + S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d + S_bar_b_j = S_b_j - jnp.mean(S_b_j, axis=2, keepdims=True) + S_bar_flat_b.append(S_bar_b_j.flatten()) + V_tilde_b_j = X_tilde_b @ tilde_W_V_b_j + V_b_j = V_tilde_b_j @ W_O_b[j] + V_flat_b.append(V_b_j.flatten()) + + # Compute cost matrix from pre-computed values + for i in range(num_heads): + for j in range(num_heads): + # Cosine similarity for S + dot_S = jnp.dot(S_bar_flat_a[i], S_bar_flat_b[j]) + norm_S_a = jnp.linalg.norm(S_bar_flat_a[i]) + norm_S_b = jnp.linalg.norm(S_bar_flat_b[j]) + cos_sim_S = dot_S / (norm_S_a * norm_S_b + epsilon) + cost_S = 1.0 - cos_sim_S + # Cosine similarity for V + dot_V = jnp.dot(V_flat_a[i], V_flat_b[j]) + norm_V_a = jnp.linalg.norm(V_flat_a[i]) + norm_V_b = jnp.linalg.norm(V_flat_b[j]) + cos_sim_V = dot_V / (norm_V_a * norm_V_b + epsilon) + cost_V = 1.0 - cos_sim_V + C[i, j] = (alpha * cost_S + (1 - alpha) * cost_V) + return C + +def compute_cost_matrix_data_independent(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, alpha=0.5): + C = np.zeros((num_heads, num_heads)) + for i in range(num_heads): + tilde_W_Q_a_i = compute_extended_weights(W_Q_a[i], b_Q_a[i]) + tilde_W_K_a_i = compute_extended_weights(W_K_a[i], b_K_a[i]) + tilde_W_V_a_i = compute_extended_weights(W_V_a[i], b_V_a[i]) + QKT_a_i = tilde_W_Q_a_i @ tilde_W_K_a_i.T + VO_a_i = tilde_W_V_a_i @ W_O_a[i] + centered_QKT_a_i = QKT_a_i - np.mean(QKT_a_i, axis=1, keepdims=True) + for j in range(num_heads): + tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j]) + tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j]) + tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j]) + QKT_b_j = tilde_W_Q_b_j @ tilde_W_K_b_j.T + VO_b_j = tilde_W_V_b_j @ W_O_b[j] + centered_QKT_b_j = QKT_b_j - np.mean(QKT_b_j, axis=1, keepdims=True) + cost = alpha * np.sum((centered_QKT_a_i - centered_QKT_b_j)**2) + (1 - alpha) * np.sum((VO_a_i - VO_b_j)**2) + C[i, j] = cost + return C + +def find_heads_permutation(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, activations_a, activations_b, alpha, data_independent): + if data_independent: + C = compute_cost_matrix_data_independent(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, alpha=alpha) + else: + C = compute_cost_matrix_presoftmax(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, activations_a, activations_b, alpha) + row_ind, col_ind = linear_sum_assignment(to_numpy(C)) + print({int(i): int(j) for i, j in zip(row_ind, col_ind)}) + return row_ind, col_ind + +# Stage 2 Function: Align Single Head (with optional optimization) +def align_single_head(W_Q_a_i, b_Q_a_i, W_K_a_i, b_K_a_i, W_V_a_i, b_V_a_i, W_O_a_i, + W_Q_b_i, b_Q_b_i, W_K_b_i, b_K_b_i, W_V_b_i, b_V_b_i, W_O_b_i, + init_method, optimize): + tilde_W_Q_a_i = compute_extended_weights(W_Q_a_i, b_Q_a_i) + tilde_W_K_a_i = compute_extended_weights(W_K_a_i, b_K_a_i) + tilde_W_V_a_i = compute_extended_weights(W_V_a_i, b_V_a_i) + Y_O_a_i = W_O_a_i.T + tilde_W_Q_b_i = compute_extended_weights(W_Q_b_i, b_Q_b_i) + tilde_W_K_b_i = compute_extended_weights(W_K_b_i, b_K_b_i) + tilde_W_V_b_i = compute_extended_weights(W_V_b_i, b_V_b_i) + Y_O_b_i = W_O_b_i.T + + if init_method == 'ortho': + A_init = solve_orthogonal(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i) + B_init = solve_orthogonal(Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i) + elif init_method == 'random': + while True: + A_init = np.random.normal(loc=1, scale=1, size=(tilde_W_Q_a_i.shape[1], tilde_W_Q_a_i.shape[1])) + if np.linalg.det(A_init) != 0: + break + while True: + B_init = np.random.normal(loc=1, scale=1, size=(tilde_W_V_a_i.shape[1], tilde_W_V_a_i.shape[1])) + if np.linalg.det(B_init) != 0: + break + elif init_method == 'identity': + A_init = np.eye(tilde_W_Q_a_i.shape[1]) + B_init = np.eye(tilde_W_V_a_i.shape[1]) + else: + raise ValueError("Invalid initialization method") + + if optimize: + A, objective_values_A, grad_norms_A, condition_nums_A = optimize_alignment( + A_init, tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i + ) + B, objective_values_B, grad_norms_B, condition_nums_B = optimize_alignment( + B_init, Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i + ) + else: + A = A_init + B = B_init + + A_inv = np.linalg.inv(A) + B_inv = np.linalg.inv(B) + W_Q_aligned = W_Q_b_i @ A.T + b_Q_aligned = b_Q_b_i @ A.T + W_K_aligned = W_K_b_i @ A_inv + b_K_aligned = b_K_b_i @ A_inv + W_V_aligned = W_V_b_i @ B_inv + b_V_aligned = b_V_b_i @ B_inv + W_O_aligned = B @ W_O_b_i + + aligned_params = { + 'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned}, + 'key': {'kernel': W_K_aligned, 'bias': b_K_aligned}, + 'value': {'kernel': W_V_aligned, 'bias': b_V_aligned}, + 'out': {'kernel': W_O_aligned} + } + + if optimize: + return { + 'aligned_params': aligned_params, + 'metrics_A': { + 'objective_values': objective_values_A, + 'grad_norms': grad_norms_A, + 'condition_nums': condition_nums_A + }, + 'metrics_B': { + 'objective_values': objective_values_B, + 'grad_norms': grad_norms_B, + 'condition_nums': condition_nums_B + } + } + return {'aligned_params': aligned_params} + +def align_attention_params_main(rng, params_a, params_b, layer_idx, num_heads, + activations_for_layer_a, activations_for_layer_b, plot_path=None, init_method='ortho', permute_heads=True, optimize=True, method_name="", alpha=0.5, data_independent=False): + + params_a_extracted, _ = extract_attention_params(params_a, layer_idx) + params_b_extracted, mha_path_b = extract_attention_params(params_b, layer_idx) + + params_a_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', + 'value', 'value_bias', 'out', 'out_bias'], params_a_extracted)} + params_b_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', + 'value', 'value_bias', 'out', 'out_bias'], params_b_extracted)} + + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_to_per_head(params_a_np, num_heads) + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_to_per_head(params_b_np, num_heads) + + if permute_heads: + row_ind, col_ind = find_heads_permutation( + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, activations_for_layer_a, activations_for_layer_b, alpha, data_independent + ) + + W_Q_b = [W_Q_b[j] for j in col_ind] + b_Q_b = [b_Q_b[j] for j in col_ind] + W_K_b = [W_K_b[j] for j in col_ind] + b_K_b = [b_K_b[j] for j in col_ind] + W_V_b = [W_V_b[j] for j in col_ind] + b_V_b = [b_V_b[j] for j in col_ind] + W_O_b = [W_O_b[j] for j in col_ind] + + if optimize: + metrics_A_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']} + metrics_B_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']} + + aligned_params = {} + for i in range(num_heads): + result = align_single_head( + W_Q_a[i], b_Q_a[i], W_K_a[i], b_K_a[i], W_V_a[i], b_V_a[i], W_O_a[i], + W_Q_b[i], b_Q_b[i], W_K_b[i], b_K_b[i], W_V_b[i], b_V_b[i], W_O_b[i], + init_method, optimize + ) + aligned_params[f'head_{i}'] = result['aligned_params'] + if optimize: + for key in metrics_A_all: + metrics_A_all[key].append(result['metrics_A'][key]) + metrics_B_all[key].append(result['metrics_B'][key]) + + query_kernel = np.stack([aligned_params[f'head_{i}']['query']['kernel'] for i in range(num_heads)], axis=1) + query_bias = np.stack([aligned_params[f'head_{i}']['query']['bias'] for i in range(num_heads)], axis=0) + key_kernel = np.stack([aligned_params[f'head_{i}']['key']['kernel'] for i in range(num_heads)], axis=1) + key_bias = np.stack([aligned_params[f'head_{i}']['key']['bias'] for i in range(num_heads)], axis=0) + value_kernel = np.stack([aligned_params[f'head_{i}']['value']['kernel'] for i in range(num_heads)], axis=1) + value_bias = np.stack([aligned_params[f'head_{i}']['value']['bias'] for i in range(num_heads)], axis=0) + out_kernel = np.stack([aligned_params[f'head_{i}']['out']['kernel'] for i in range(num_heads)], axis=0) + + return_dict = { + 'aligned_params': { + 'query': {'kernel': jnp.array(query_kernel), 'bias': jnp.array(query_bias)}, + 'key': {'kernel': jnp.array(key_kernel), 'bias': jnp.array(key_bias)}, + 'value': {'kernel': jnp.array(value_kernel), 'bias': jnp.array(value_bias)}, + 'out': {'kernel': jnp.array(out_kernel), 'bias': params_b_np['out_bias']} + }, + 'mha_path': mha_path_b # Return the path for updating model_b + } + if optimize: + return_dict['metrics_A_all'] = metrics_A_all + return_dict['metrics_B_all'] = metrics_B_all + return return_dict + +def matching_attn(rng, params_a, params_b, activations_a, activations_b, finetune_layer_which, num_heads, plot_path): + params_dict = {} + configurations = [ + ("data_indep_permu_head_init_ortho_no_opt", True, 'ortho', True, False), + #("data_dep_permu_head_init_ortho_no_opt", False, 'ortho', True, False), + #("data_dep_permu_head_init_ortho_opt", 'ortho', True, True), + ] + + for name, data_independent, init_method, permute_heads, optimize in configurations: + aligned_params = copy.deepcopy(params_b) + if optimize: + layer_to_metrics_A = {} + layer_to_metrics_B = {} + for layer_idx in finetune_layer_which: + activations_for_layer_a = activations_a[layer_idx] + activations_for_layer_b = activations_b[layer_idx] + + result = align_attention_params_main( + rng, params_a, aligned_params, layer_idx, num_heads, activations_for_layer_a, activations_for_layer_b, plot_path=None, + init_method=init_method, permute_heads=permute_heads, optimize=optimize, method_name=name, data_independent=data_independent + ) + + # Update the aligned_params tree using the path returned by the alignment function. + unfrozen_params = unfreeze(aligned_params) + + # Navigate to the parent dictionary of the MHA block + temp_dict = unfrozen_params + for key in result['mha_path'][:-1]: + temp_dict = temp_dict[key] + + # Update the MHA block with the aligned parameters + temp_dict[result['mha_path'][-1]] = result['aligned_params'] + + aligned_params = freeze(unfrozen_params) + + if optimize: + layer_to_metrics_A[layer_idx] = result['metrics_A_all'] + layer_to_metrics_B[layer_idx] = result['metrics_B_all'] + + total_sum = tree_util.tree_reduce(lambda acc, x: acc + jnp.sum(x), aligned_params, initializer=0) + print(f"{name}: {total_sum}, sanity check") + params_dict[name] = aligned_params + + + if optimize and plot_path: + os.makedirs(plot_path, exist_ok=True) + num_layers = len(finetune_layer_which) + layers = finetune_layer_which + metric_keys = ['objective_values', 'grad_norms', 'condition_nums'] + for metric_key in metric_keys: + fig, axs = plt.subplots(num_layers + 1, 2, figsize=(20, 5 * (num_layers + 1)), sharex='col') + for col in range(2): + if col == 0: + metrics_per_layer = layer_to_metrics_A + alignment_type = "Query/Key Alignment" + else: + metrics_per_layer = layer_to_metrics_B + alignment_type = "Value/Out Alignment" + + # Plot per-layer subplots + for row in range(num_layers): + layer = layers[row] + data_list = metrics_per_layer[layer][metric_key] + labels = [f"Head {i}" for i in range(num_heads)] + ax = axs[row, col] + for data, label in zip(data_list, labels): + if data: + final_val = data[-1] + ax.plot(data, label=f"{label} ({final_val:.4f})") + ax.set_title(f"Layer {layer}: {metric_key.replace('_', ' ').capitalize()} - {alignment_type}") + ax.set_xlabel('Iteration') + ax.set_ylabel(metric_key.replace('_', ' ').capitalize()) + ax.legend(loc='center left', bbox_to_anchor=(1, 0.5)) + + # Bottom row: mean across heads for all layers + ax = axs[num_layers, col] + data_list = [] + labels = [] + for layer in layers: + head_data = metrics_per_layer[layer][metric_key] + if head_data: + max_len = max(len(d) for d in head_data if d) + padded = [] + for d in head_data: + if d: + if len(d) < max_len: + last = d[-1] + padded.append(d + [last] * (max_len - len(d))) + else: + padded.append(d) + if padded: + mean_data = np.mean(padded, axis=0).tolist() + data_list.append(mean_data) + final_mean = mean_data[-1] + labels.append(f"Layer {layer} ({final_mean:.4f})") + for data, label in zip(data_list, labels): + ax.plot(data, label=label) + ax.set_title(f"All Layers Mean: {metric_key.replace('_', ' ').capitalize()} - {alignment_type}") + ax.set_xlabel('Iteration') + ax.set_ylabel(metric_key.replace('_', ' ').capitalize()) + ax.legend(loc='center left', bbox_to_anchor=(1, 0.5)) + + plt.tight_layout() + save_path = os.path.join(plot_path, f"{name}_{metric_key}.png") + plt.savefig(save_path, bbox_inches='tight') + plt.close() + + return params_dict + +#############RoPE################# +def get_rope_matrix(seq_len, d_head): + """Generates RoPE rotation matrices R[m] of shape (seq_len, d_head/2, 2, 2).""" + assert d_head % 2 == 0, "d_head must be even" + # inv_freq: (d_head/2,) + inv_freq = 1.0 / (10000 ** (jnp.arange(0, d_head, 2) / d_head)) + t = jnp.arange(seq_len) # (seq_len,) + freqs = jnp.einsum('i,j->ij', t, inv_freq) # (seq_len, d_head/2) + + cos_freqs = jnp.cos(freqs) # (seq_len, h) + sin_freqs = jnp.sin(freqs) # (seq_len, h) + + # Build rotation matrices per position and subspace: shape (seq_len, h, 2, 2) + # Each 2x2 is [[cos, -sin], [sin, cos]] + R = jnp.stack( + [ + jnp.stack([cos_freqs, -sin_freqs], axis=-1), # (seq_len, h, 2) -> first row entries + jnp.stack([sin_freqs, cos_freqs], axis=-1), # (seq_len, h, 2) -> second row entries + ], + axis=-2 + ) # After this stack: shape (seq_len, h, 2, 2) + return R + +@jax.jit +def apply_rope(x, R): + """Applies RoPE to x of shape (B, L, D_k) using R shape (L, D_k/2, 2, 2).""" + B, L, Dk = x.shape + assert Dk % 2 == 0 + x_pairs = x.reshape((B, L, Dk//2, 2)) # (B, L, h, 2) + # R must be (L, h, 2, 2) + # einsum: 'b l h c, l h c r -> b l h r' -> back to (B, L, h, 2) + x_rotated = jnp.einsum('blhc,lhcr->blhr', x_pairs, R) + return x_rotated.reshape((B, L, Dk)) + +def compute_cost_matrix_presoftmax_rope(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, activations_a, activations_b, alpha=0.5, epsilon=1e-8): + B_a, L_a, D_a = activations_a.shape + B_b, L_b, D_b = activations_b.shape + assert B_a == B_b and L_a == L_b and D_a == D_b + L, d_head = L_a, W_Q_a[0].shape[1] + + rope_matrices = get_rope_matrix(L, d_head) + ones_col_a, ones_col_b = jnp.ones((B_a, L_a, 1)), jnp.ones((B_b, L_b, 1)) + X_tilde_a = jnp.concatenate([activations_a, ones_col_a], axis=-1) + X_tilde_b = jnp.concatenate([activations_b, ones_col_b], axis=-1) + sqrt_d = jnp.sqrt(float(d_head)) + C = np.zeros((num_heads, num_heads)) + + # Pre-compute for model A + S_bar_flat_a, V_flat_a = [], [] + for i in range(num_heads): + tilde_W_Q_a_i, tilde_W_K_a_i, tilde_W_V_a_i = compute_extended_weights(W_Q_a[i], b_Q_a[i]), compute_extended_weights(W_K_a[i], b_K_a[i]), compute_extended_weights(W_V_a[i], b_V_a[i]) + Q_rope_a_i, K_rope_a_i = apply_rope(X_tilde_a @ tilde_W_Q_a_i, rope_matrices), apply_rope(X_tilde_a @ tilde_W_K_a_i, rope_matrices) + S_a_i = jnp.einsum('bld,bmd->blm', Q_rope_a_i, K_rope_a_i) / sqrt_d + S_bar_flat_a.append((S_a_i - jnp.mean(S_a_i, axis=2, keepdims=True)).flatten()) + V_flat_a.append((X_tilde_a @ tilde_W_V_a_i @ W_O_a[i]).flatten()) + + # Pre-compute for model B + S_bar_flat_b, V_flat_b = [], [] + for j in range(num_heads): + tilde_W_Q_b_j, tilde_W_K_b_j, tilde_W_V_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j]), compute_extended_weights(W_K_b[j], b_K_b[j]), compute_extended_weights(W_V_b[j], b_V_b[j]) + Q_rope_b_j, K_rope_b_j = apply_rope(X_tilde_b @ tilde_W_Q_b_j, rope_matrices), apply_rope(X_tilde_b @ tilde_W_K_b_j, rope_matrices) + S_b_j = jnp.einsum('bld,bmd->blm', Q_rope_b_j, K_rope_b_j) / sqrt_d + S_bar_flat_b.append((S_b_j - jnp.mean(S_b_j, axis=2, keepdims=True)).flatten()) + V_flat_b.append((X_tilde_b @ tilde_W_V_b_j @ W_O_b[j]).flatten()) + + for i in range(num_heads): + for j in range(num_heads): + dot_S = jnp.dot(S_bar_flat_a[i], S_bar_flat_b[j]) + norm_S_a, norm_S_b = jnp.linalg.norm(S_bar_flat_a[i]), jnp.linalg.norm(S_bar_flat_b[j]) + cost_S = 1.0 - (dot_S / (norm_S_a * norm_S_b + epsilon)) + dot_V = jnp.dot(V_flat_a[i], V_flat_b[j]) + norm_V_a, norm_V_b = jnp.linalg.norm(V_flat_a[i]), jnp.linalg.norm(V_flat_b[j]) + cost_V = 1.0 - (dot_V / (norm_V_a * norm_V_b + epsilon)) + C[i, j] = (alpha * cost_S + (1 - alpha) * cost_V) + return C + +def find_heads_permutation_rope(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, activations_a, activations_b, alpha, data_independent): + """Wrapper to find permutation using the RoPE cost matrix with model-specific activations.""" + if data_independent: + C = compute_cost_matrix_data_independent(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, alpha=alpha) + else: + C = compute_cost_matrix_presoftmax_rope( + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, activations_a, activations_b, alpha + ) + row_ind, col_ind = linear_sum_assignment(to_numpy(C)) + print("RoPE Head Permutation:", {int(i): int(j) for i, j in zip(row_ind, col_ind)}) + return row_ind, col_ind + +from scipy.optimize import minimize_scalar +from scipy.linalg import block_diag +from math import sqrt, cos, sin, atan2 +import numpy as np + +def solve_rope_qk_alignment(W_Q_a_i, b_Q_a_i, W_K_a_i, b_K_a_i, + W_Q_b_i, b_Q_b_i, W_K_b_i, b_K_b_i): + """ + Solves for the G_RoPE alignment matrix U for a single head's QK weights. + """ + tilde_W_Q_a = compute_extended_weights(W_Q_a_i, b_Q_a_i) + tilde_W_K_a = compute_extended_weights(W_K_a_i, b_K_a_i) + tilde_W_Q_b = compute_extended_weights(W_Q_b_i, b_Q_b_i) + tilde_W_K_b = compute_extended_weights(W_K_b_i, b_K_b_i) + + D_k = tilde_W_Q_a.shape[1] + assert D_k % 2 == 0, "Head dimension must be even for RoPE." + + U_blocks = [] + J = jnp.array([[0, -1], [1, 0]]) + + for j in range(D_k // 2): + # 1. Slice submatrices for the j-th 2D subspace + sl = slice(2 * j, 2 * j + 2) + Q_a_j, Q_b_j = tilde_W_Q_a[:, sl], tilde_W_Q_b[:, sl] + K_a_j, K_b_j = tilde_W_K_a[:, sl], tilde_W_K_b[:, sl] + + # 2. Precompute constants + N_Q = jnp.sum(Q_b_j**2) + N_K = jnp.sum(K_b_j**2) + C_Q = Q_a_j.T @ Q_b_j + C_K = K_a_j.T @ K_b_j + + c_q = 0.5 * (jnp.trace(C_Q) + 1j * jnp.trace(C_Q @ J)) + c_k = 0.5 * (jnp.trace(C_K) + 1j * jnp.trace(C_K @ J)) + + A = jnp.abs(c_q)**2 + B = jnp.abs(c_k)**2 + C = 2 * jnp.real(c_q * jnp.conj(c_k)) + + # Convert all JAX/Device arrays to native Python/NumPy types for SciPy + N_Q_f, N_K_f = float(N_Q), float(N_K) + A_f, B_f, C_f = float(A), float(B), float(C) + c_q_f = complex(c_q) + c_k_f = complex(c_k) + + # Define the 1D scalar objective function using native floats + def g_objective(x): + x = float(x) + # Protect the sqrt argument from tiny negative values due to roundoff + inner_term = A_f * x + (B_f / x) + C_f + safe_inner = max(inner_term, 1e-20) + return x * N_Q_f + N_K_f / x - 4.0 * sqrt(safe_inner) + + # 4. Find the minimizer x* using robust bounds + res = minimize_scalar(g_objective, bounds=(1e-8, 1e8), method='bounded') + x_star = res.x + + # 5. Reconstruct the optimal 2x2 alignment matrix U_j using NumPy/math + r_star = sqrt(x_star) + combined_c = r_star * c_q_f + (1 / r_star) * c_k_f + if abs(combined_c) < 1e-30: + theta_star = 0.0 + else: + theta_star = -atan2(combined_c.imag, combined_c.real) + a = r_star * cos(theta_star) + b = r_star * sin(theta_star) + + U_j = np.array([[a, -b], [b, a]]) + U_blocks.append(U_j) + + # 6. Assemble the full block-diagonal matrix U + U_opt = block_diag(*U_blocks) + condU = np.linalg.cond(U_opt) + if condU > 1e12: + # fallback: scale blocks to have minimum magnitude, or add small diag: + eps = 1e-6 + U_opt = U_opt + eps * np.eye(U_opt.shape[0]) + return U_opt + +def align_attention_params_main_rope(params_a, params_b, layer_idx, num_heads, + activations_for_layer_a, activations_for_layer_b, + init_method_vo='ortho', permu_heads=True, optimize_vo=True, alpha=0.5, data_independent=False): + """ + Aligns a single MHA layer with RoPE using model-specific activations. + COMPATIBLE with multiple model structures. + """ + # Use the compatible extractor to get params and the update path --- + params_a_extracted, _ = extract_attention_params(params_a, layer_idx) + params_b_extracted, mha_path_b = extract_attention_params(params_b, layer_idx) + params_a_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', 'value', 'value_bias', 'out', 'out_bias'], params_a_extracted)} + params_b_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', 'value', 'value_bias', 'out', 'out_bias'], params_b_extracted)} + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_to_per_head(params_a_np, num_heads) + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_to_per_head(params_b_np, num_heads) + + if permu_heads: + # --- Stage 1: Head Permutation (RoPE version) --- + row_ind, col_ind = find_heads_permutation_rope( + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, activations_for_layer_a, activations_for_layer_b, alpha, data_independent + ) + # --- Reorder heads of model B --- (No changes here) + W_Q_b = [W_Q_b[j] for j in col_ind] + b_Q_b = [b_Q_b[j] for j in col_ind] + W_K_b = [W_K_b[j] for j in col_ind] + b_K_b = [b_K_b[j] for j in col_ind] + W_V_b = [W_V_b[j] for j in col_ind] + b_V_b = [b_V_b[j] for j in col_ind] + W_O_b = [W_O_b[j] for j in col_ind] + + # --- Stage 2: Per-Head Parameter Alignment --- (No changes here) + aligned_params_list = [] + for i in range(num_heads): + # ... (rest of the function is identical) + # QK Alignment (RoPE specific) + U = solve_rope_qk_alignment( + W_Q_a[i], b_Q_a[i], W_K_a[i], b_K_a[i], + W_Q_b[i], b_Q_b[i], W_K_b[i], b_K_b[i] + ) + U_inv = np.linalg.inv(U) + + W_Q_aligned = W_Q_b[i] @ U.T + b_Q_aligned = b_Q_b[i] @ U.T + W_K_aligned = W_K_b[i] @ U_inv + b_K_aligned = b_K_b[i] @ U_inv + + # VO Alignment (Standard MHA logic, as it's unaffected by RoPE) + tilde_W_V_a_i = compute_extended_weights(W_V_a[i], b_V_a[i]) + Y_O_a_i = W_O_a[i].T + tilde_W_V_b_i = compute_extended_weights(W_V_b[i], b_V_b[i]) + Y_O_b_i = W_O_b[i].T + + if init_method_vo == 'ortho': + B_init = solve_orthogonal(Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i) + else: + B_init = np.identity(W_O_b[i].shape[0]) + + if optimize_vo: + B, _, _, _ = optimize_alignment(B_init, Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i) + else: + B = B_init + + B_inv = np.linalg.inv(B) + W_V_aligned = W_V_b[i] @ B_inv + b_V_aligned = b_V_b[i] @ B_inv + W_O_aligned = B @ W_O_b[i] + + aligned_params_list.append({ + 'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned}, + 'key': {'kernel': W_K_aligned, 'bias': b_K_aligned}, + 'value': {'kernel': W_V_aligned, 'bias': b_V_aligned}, + 'out': {'kernel': W_O_aligned} + }) + + # --- Reassemble Parameters --- (No changes here) + query_kernel = np.stack([p['query']['kernel'] for p in aligned_params_list], axis=1) + query_bias = np.stack([p['query']['bias'] for p in aligned_params_list], axis=0) + key_kernel = np.stack([p['key']['kernel'] for p in aligned_params_list], axis=1) + key_bias = np.stack([p['key']['bias'] for p in aligned_params_list], axis=0) + value_kernel = np.stack([p['value']['kernel'] for p in aligned_params_list], axis=1) + value_bias = np.stack([p['value']['bias'] for p in aligned_params_list], axis=0) + out_kernel = np.stack([p['out']['kernel'] for p in aligned_params_list], axis=0) + + return { + 'aligned_params': { + 'query': {'kernel': jnp.array(query_kernel), 'bias': jnp.array(query_bias)}, + 'key': {'kernel': jnp.array(key_kernel), 'bias': jnp.array(key_bias)}, + 'value': {'kernel': jnp.array(value_kernel), 'bias': jnp.array(value_bias)}, + 'out': {'kernel': jnp.array(out_kernel), 'bias': jnp.array(params_b_np['out_bias'])} + }, + 'mha_path': mha_path_b + } + +def matching_attn_rope(params_a, params_b, activations_a, activations_b, finetune_layer_which, num_heads, alpha=0.5): + """ + Main function to align RoPE-based MHA layers using separate activations for each model. + COMPATIBLE with multiple model structures. + + Args: + params_a (dict): Parameters of the reference model A. + params_b (dict): Parameters of the model B to be aligned. + activations (dict): Dictionary mapping layer_idx to input activations. + finetune_layer_which (list): List of layer indices to align. + num_heads (int): Number of attention heads. + alpha (float): Weighting factor for permutation cost matrix calculation. + + Returns: + dict: A dictionary where keys are method names and values are the + aligned parameters for model B. + """ + params_dict = {} + + # Define a single configuration for RoPE alignment. + # The VO part can optionally be optimized after orthogonal initialization. + configurations = [ + ("data_indep_permu_head_init_ortho_no_opt", True, 'ortho', True, False), + #("data_dep_permu_head_init_ortho_no_opt", False, 'ortho', True, False), + #("data_dep_permu_head_init_ortho_opt", 'ortho', True, True), + ] + + for name, data_independent, init_method_vo, permu_heads, optimize_vo in configurations: + print(f"--- Running RoPE Alignment Configuration: {name} ---") + + # Start with a deepcopy of params_b to modify iteratively --- + aligned_params = copy.deepcopy(params_b) + + for layer_idx in finetune_layer_which: + print(f"Aligning Layer {layer_idx}...") + + activations_for_layer_a = activations_a[layer_idx] + activations_for_layer_b = activations_b[layer_idx] + + # Call the main alignment function for a single RoPE MHA layer + result = align_attention_params_main_rope( + params_a=params_a, + params_b=aligned_params, + layer_idx=layer_idx, + num_heads=num_heads, + activations_for_layer_a=activations_for_layer_a, + activations_for_layer_b=activations_for_layer_b, + init_method_vo=init_method_vo, + permu_heads=permu_heads, + optimize_vo=optimize_vo, + alpha=alpha, + data_independent=data_independent + ) + + # Update the aligned_params tree using the path --- + unfrozen_params = unfreeze(aligned_params) + + # Navigate to the parent dictionary of the MHA block + temp_dict = unfrozen_params + for key in result['mha_path'][:-1]: + temp_dict = temp_dict[key] + + # Update the MHA block with the aligned parameters + temp_dict[result['mha_path'][-1]] = result['aligned_params'] + + aligned_params = freeze(unfrozen_params) + + + # Sanity check + total_sum = tree_util.tree_reduce(lambda acc, x: acc + jnp.sum(x), aligned_params, initializer=0) + print(f"Finished configuration '{name}'. Total parameter sum: {total_sum:.4f}\n") + + params_dict[name] = aligned_params + + return params_dict + + +# this function performs matching on all possible heads permutations +# thus only support len(finetune_layer_which)==1 i.e. at 1 layer only +import itertools + +layer_key_prefix = 'TransformerEncoderLayer' +attention_key = 'MultiHeadDotProductAttention_0' + +def matching_attn_all_heads_permu(rng, params_a, params_b, finetune_layer_which, num_heads, plot_path, rope_use=False, activations_a=None, activations_b=None): + """ + Performs matching for all possible head permutations for a single layer and records the + optimal permutation for several data-dependent and independent methods. + + Args: + rng: JAX random key. + params_a: Parameters of the first model. + params_b: Parameters of the second model. + finetune_layer_which: A list containing the index of the layer to finetune (must have length 1). + num_heads: The number of attention heads. + plot_path: Path for saving plots (not used in this version but kept for consistency). + rope_use: Boolean indicating if RoPE is used in the model. + activations_a: A dictionary of activations from model A, keyed by layer index. + activations_b: A dictionary of activations from model B, keyed by layer index. + + Returns: + A tuple containing: + - params_dict: Dictionary of aligned parameters for different settings and permutations. + - heads_objective_values: Dictionary of objective values for each permutation. + - heads_permutation_sol: Dictionary storing the optimal permutation for each calculation method. + """ + assert len(finetune_layer_which) == 1, "This function only supports one layer at a time." + + params_dict = defaultdict(lambda: defaultdict(dict)) + heads_objective_values = defaultdict(lambda: defaultdict(dict)) + heads_permutation_sol = {} + C_dict = {} + + layer_idx = finetune_layer_which[0] + layer_key = f'{layer_key_prefix}_{layer_idx}' + + # --- Extract and reshape weights --- + params_a_extracted = extract_attention_params(params_a, layer_key, attention_key) + params_b_extracted = extract_attention_params(params_b, layer_key, attention_key) + params_a_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', 'value', 'value_bias', 'out', 'out_bias'], params_a_extracted)} + params_b_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', 'value', 'value_bias', 'out', 'out_bias'], params_b_extracted)} + W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_to_per_head(params_a_np, num_heads) + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_to_per_head(params_b_np, num_heads) + + # --- Define permutation settings --- + head_permu_settings = [ + ("data-independent", True, None), + #("data-dependent_acti-b-use", False, True), + #("data-dependent_acti-b-notuse", False, False), + ] + + alpha = 0.5 # Using a fixed alpha as in the original script + + # --- Calculate cost matrices and find optimal permutations for each method --- + print("Calculating cost matrices and optimal permutations for each method...") + for name, data_independent, activation_b_use in head_permu_settings: + if data_independent: + C = compute_cost_matrix_data_independent(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, alpha=alpha) + else: + # For data-dependent methods + assert activations_a is not None, "Activations for model A must be provided for data-dependent methods." + # Use model B's activations if specified, otherwise use model A's activations for both + current_activations_b = activations_b[layer_idx] if activation_b_use and activations_b else activations_a[layer_idx] + + if rope_use: + C = compute_cost_matrix_presoftmax_rope(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, activations_a[layer_idx], current_activations_b, alpha=alpha) + else: + C = compute_cost_matrix_presoftmax(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, + W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, + num_heads, activations_a[layer_idx], current_activations_b, alpha=alpha) + + C_dict[name] = C + row_ind, col_ind = linear_sum_assignment(C) + # Store the permutation mapping row_ind (model A) to col_ind (model B) + # We sort by row_ind to ensure the permutation is always in the order [perm_for_head_0, perm_for_head_1, ...] + permutation_solution = [int(col_ind[i]) for i in np.argsort(row_ind)] + heads_permutation_sol[name] = ([int(_) for _ in row_ind], [int(_) for _ in col_ind]) + print(f" - Method '{name}': Optimal permutation is {permutation_solution}") + + + # --- Iterate through all possible permutations for alignment and interpolation --- + permutations_to_evaluate = list(itertools.permutations(range(num_heads))) + + # Collect unique permutations from heads_permutation_sol + S = set() + for name, _, _ in head_permu_settings: + _, col_ind = heads_permutation_sol[name] + col_ind = [int(_) for _ in col_ind] + S.add(tuple(col_ind)) # Convert to tuple for set compatibility + + # Initialize permutations to use + permutations_to_use = set(S) + max_permutations = 24 + + # Sample additional permutations up to max_permutations + while len(permutations_to_use) < max_permutations: + perm = tuple(np.random.permutation(num_heads)) + permutations_to_use.add(perm) # Set ensures no duplicates + + permutations_to_evaluate = permutations_to_use + print(f"\nEvaluating all {len(permutations_to_evaluate)} possible head permutations...") + + # Define alignment settings + alignment_settings = [ + ("init_ortho_no_opt", 'ortho', False), + #("init_ortho_opt", 'ortho', True), + ] + + for setting_name, init_method, optimize in alignment_settings: + for perm_tuple in permutations_to_evaluate: + perm = list(perm_tuple) + # Create a fresh copy of model B parameters for this specific permutation + params_b_permuted = copy.deepcopy(unfreeze(params_b)) + + # Manually permute the attention heads of model B for the specified layer + # This is done by re-indexing the weight tensors according to `perm` + params_b_layer = params_b_permuted[layer_key][attention_key] + perm_array = np.array(perm) + + # Permute dimensions related to heads + params_b_layer['query']['kernel'] = params_b_layer['query']['kernel'][:, perm_array, :] + params_b_layer['query']['bias'] = params_b_layer['query']['bias'][perm_array, :] + params_b_layer['key']['kernel'] = params_b_layer['key']['kernel'][:, perm_array, :] + params_b_layer['key']['bias'] = params_b_layer['key']['bias'][perm_array, :] + params_b_layer['value']['kernel'] = params_b_layer['value']['kernel'][:, perm_array, :] + params_b_layer['value']['bias'] = params_b_layer['value']['bias'][perm_array, :] + params_b_layer['out']['kernel'] = params_b_layer['out']['kernel'][perm_array, :, :] + + # Align the permuted model B to model A + if rope_use: + aligned_result = align_attention_params_main_rope( + params_a, params_b_permuted, layer_idx, num_heads, + activations_a[layer_idx], + activations_b[layer_idx], + init_method_vo=init_method, permu_heads=False, optimize_vo=optimize, alpha=alpha + ) + # The rope main function returns the aligned MHA block directly + result_to_store = aligned_result + else: + # Standard alignment. `permute_heads` is False because we do it manually above. + result_dict = align_attention_params_main( + rng, params_a, params_b_permuted, layer_idx, num_heads, + activations_a[layer_idx], + activations_b[layer_idx], + plot_path=None, init_method=init_method, + permute_heads=False, optimize=optimize, alpha=alpha + ) + result_to_store = result_dict + + perm_str = str([int(_) for _ in perm]) + params_dict[setting_name][perm_str] = result_to_store + + # --- Calculate objective cost for this permutation against each method's cost matrix --- + for method_name, C_matrix in C_dict.items(): + # The cost is the sum of C[i, j] for the mapping i -> perm[i] + total_cost = sum(C_matrix[i, perm[i]] for i in range(num_heads)) + heads_objective_values[method_name][setting_name][perm_str] = float(total_cost) + + print("Finished evaluating all permutations.") + return params_dict, heads_objective_values, heads_permutation_sol + +##########Transformers matching +import copy +import jax.numpy as jnp +import numpy as np +from flax.core import unfreeze, freeze +from scipy.optimize import linear_sum_assignment + +# --- Helper Functions (from your draft and my previous code) --- + +def get_nested_item(d, keys): + """Accesses a nested dictionary item using a tuple of keys.""" + for key in keys: + d = d[key] + return d + +def set_nested_item(d, keys, value): + """Sets a value in a nested dictionary using a tuple of keys.""" + current = d + for key in keys[:-1]: + current = current[key] + current[keys[-1]] = value + +def extract_ffn_params(params, layer_idx): + """ + Extracts FFN parameters from different, known model structures in a compatible way. + This function detects the model type and constructs the correct path to the + FFN parameters for a given layer index. + Args: + params: The Flax parameter tree. + layer_idx: The integer index of the transformer layer. + Returns: + A tuple containing: + - W1, b1, W2 (FFN weights and biases for Dense_0 and Dense_1). + - dense0_path, dense1_path (tuples representing the nested paths to Dense_0 and Dense_1). + """ + # Detect cifar_vit style: params['TransformerEncoderLayer_...'] + if f'TransformerEncoderLayer_{layer_idx}' in params: + base_path_tuple = (f'TransformerEncoderLayer_{layer_idx}',) + dense0_path = base_path_tuple + ('Dense_0',) + dense1_path = base_path_tuple + ('Dense_1',) + # Detect vit-jax style: params['Transformer']['encoderblock_...'] + elif 'Transformer' in params and f'encoderblock_{layer_idx}' in params['Transformer']: + # Note: vit-jax often nests the MLP in its own block, e.g., 'MlpBlock_0' + # We check for its existence for robustness. + encoder_block = get_nested_item(params, ('Transformer', f'encoderblock_{layer_idx}')) + mlp_key = next((k for k in encoder_block if 'MlpBlock' in k), None) + if mlp_key: + base_path_tuple = ('Transformer', f'encoderblock_{layer_idx}', mlp_key) + else: # Fallback if no explicit MlpBlock + base_path_tuple = ('Transformer', f'encoderblock_{layer_idx}') + dense0_path = base_path_tuple + ('Dense_0',) + dense1_path = base_path_tuple + ('Dense_1',) + else: + raise KeyError(f"Could not find a known path for FFN layer {layer_idx} in the provided params.") + + dense0_block = get_nested_item(params, dense0_path) + dense1_block = get_nested_item(params, dense1_path) + + W1 = dense0_block['kernel'] + b1 = dense0_block['bias'] + W2 = dense1_block['kernel'] + + return W1, b1, W2, dense0_path, dense1_path + +# --- Core Functions --- + +def matching_transformer_ffn(params_a, params_b, finetune_layer_which): + """ + Aligns the FFN components of two models for the specified layers. + + This function computes the optimal permutation of hidden neurons in the FFN + of model B to match model A. The permutation is found by solving a Linear + Assignment Problem (LAP) where the cost is the sum of squared L2 distances + between the incoming (weights + bias) and outgoing weights of each neuron pair. + + Args: + params_a (dict): Parameters of the reference model A. + params_b (dict): Parameters of the model B to be aligned. + finetune_layer_which (list): List of layer indices to align. + + Returns: + dict: The aligned parameters for model B. + """ + aligned_params_b = copy.deepcopy(params_b) + + for layer_idx in finetune_layer_which: + # Extract parameters and paths for both models + W1_a, b1_a, W2_a, _, _ = extract_ffn_params(params_a, layer_idx) + W1_b, b1_b, W2_b, dense0_path_b, dense1_path_b = extract_ffn_params(aligned_params_b, layer_idx) + + # Convert to NumPy for computation + W1_a, b1_a, W2_a = np.array(W1_a), np.array(b1_a), np.array(W2_a) + W1_b, b1_b, W2_b = np.array(W1_b), np.array(b1_b), np.array(W2_b) + + D_hidden = W1_a.shape[1] + assert W1_b.shape[1] == D_hidden, f"FFN hidden dimensions for layer {layer_idx} must match." + + # Compute cost matrix C + C = np.zeros((D_hidden, D_hidden), dtype=np.float32) + for i in range(D_hidden): + # Incoming weights and bias for neuron i of model A + in_a = np.concatenate([W1_a[:, i], [b1_a[i]]]) + # Outgoing weights for neuron i of model A + out_a = W2_a[i, :] + for j in range(D_hidden): + in_b = np.concatenate([W1_b[:, j], [b1_b[j]]]) + out_b = W2_b[j, :] + + # Cost is the sum of squared Euclidean distances + cost = np.linalg.norm(in_a - in_b)**2 + np.linalg.norm(out_a - out_b)**2 + C[i, j] = cost + + # Solve LAP. `col_ind` gives the permutation for model B's neurons. + row_ind, col_ind = linear_sum_assignment(C) + + # Permute the weights of model B according to the solution + W1_aligned = W1_b[:, col_ind] + b1_aligned = b1_b[col_ind] + W2_aligned = W2_b[col_ind, :] + + # Update the parameter dictionary for model B + unfrozen_params = unfreeze(aligned_params_b) + set_nested_item(unfrozen_params, dense0_path_b + ('kernel',), jnp.array(W1_aligned)) + set_nested_item(unfrozen_params, dense0_path_b + ('bias',), jnp.array(b1_aligned)) + set_nested_item(unfrozen_params, dense1_path_b + ('kernel',), jnp.array(W2_aligned)) + aligned_params_b = freeze(unfrozen_params) + + return aligned_params_b + +def matching_transformer_block(rng, params_a, params_b, activations_a, activations_b, finetune_layer_which, num_heads, rope_use=False, plot_path=None): + """ + Aligns entire Transformer blocks by sequentially aligning their MHA and FFN components. + + This function first calls the appropriate MHA alignment function (`matching_attn` or + `matching_attn_rope`), which returns a dictionary of aligned parameters for various + configurations. It then iterates through this dictionary, applying the FFN + alignment to each MHA-aligned model. + + Args: + rng: JAX random key. + params_a (dict): Parameters of the reference model A. + params_b (dict): Parameters of the model B to be aligned. + activations_a (dict): Dictionary of activations from model A, keyed by layer index. + activations_b (dict): Dictionary of activations from model B, keyed by layer index. + finetune_layer_which (list): List of layer indices to align. + num_heads (int): Number of attention heads. + rope_use (bool): If True, use RoPE-specific MHA alignment. + plot_path (str, optional): Path for saving diagnostic plots. + + Returns: + dict: A dictionary where keys are configuration names (e.g., 'data_indep_...`) + and values are the fully aligned (MHA + FFN) parameter dictionaries. + """ + print("--- Starting Transformer Block Alignment ---") + + # Step 1: Align the MHA component for all specified layers and configurations. + print("\nStep 1: Aligning Multi-Head Attention components...") + if rope_use: + mha_aligned_params_dict = matching_attn_rope(params_a, params_b, activations_a, activations_b, finetune_layer_which, num_heads) + else: + mha_aligned_params_dict = matching_attn(rng, params_a, params_b, activations_a, activations_b, finetune_layer_which, num_heads, plot_path) + print("MHA alignment complete.") + + # Step 2: For each MHA-aligned model, align its FFN component. + print("\nStep 2: Aligning Feed-Forward Network components for each configuration...") + fully_aligned_params_dict = {} + for config_name, mha_aligned_params in mha_aligned_params_dict.items(): + print(f" - Aligning FFN for configuration: '{config_name}'") + fully_aligned_params = matching_transformer_ffn(params_a, mha_aligned_params, finetune_layer_which) + fully_aligned_params_dict[config_name] = fully_aligned_params + + print("FFN alignment complete.") + print("\n--- Transformer Block Alignment Finished ---") + + return fully_aligned_params_dict \ No newline at end of file diff --git a/src/online_stats.py b/src/online_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..06a7b69d9a4c20bba7f17b423be91b6b68712f80 --- /dev/null +++ b/src/online_stats.py @@ -0,0 +1,89 @@ +"""Online-ish Pearson correlation of all n x n variable pairs simultaneously.""" +from typing import NamedTuple + +import jax.numpy as jnp + + +class OnlineMean(NamedTuple): + sum: jnp.ndarray + count: int + + @staticmethod + def init(num_features: int): + return OnlineMean(sum=jnp.zeros(num_features), count=0) + + def update(self, batch: jnp.ndarray): + return OnlineMean(self.sum + jnp.sum(batch, axis=0), self.count + batch.shape[0]) + + def mean(self): + return self.sum / self.count + +class OnlineCovariance(NamedTuple): + a_mean: jnp.ndarray # (d, ) + b_mean: jnp.ndarray # (d, ) + cov: jnp.ndarray # (d, d) + var_a: jnp.ndarray # (d, ) + var_b: jnp.ndarray # (d, ) + count: int + + @staticmethod + def init(a_mean: jnp.ndarray, b_mean: jnp.ndarray): + assert a_mean.shape == b_mean.shape + assert len(a_mean.shape) == 1 + d = a_mean.shape[0] + return OnlineCovariance(a_mean, + b_mean, + cov=jnp.zeros((d, d)), + var_a=jnp.zeros((d, )), + var_b=jnp.zeros((d, )), + count=0) + + def update(self, a_batch, b_batch): + assert a_batch.shape == b_batch.shape + batch_size, _ = a_batch.shape + a_res = a_batch - self.a_mean + b_res = b_batch - self.b_mean + return OnlineCovariance(a_mean=self.a_mean, + b_mean=self.b_mean, + cov=self.cov + a_res.T @ b_res, + var_a=self.var_a + jnp.sum(a_res**2, axis=0), + var_b=self.var_b + jnp.sum(b_res**2, axis=0), + count=self.count + batch_size) + + def covariance(self): + return self.cov / (self.count - 1) + + def a_variance(self): + return self.var_a / (self.count - 1) + + def b_variance(self): + return self.var_b / (self.count - 1) + + def a_stddev(self): + return jnp.sqrt(self.a_variance()) + + def b_stddev(self): + return jnp.sqrt(self.b_variance()) + + def E_ab(self): + return self.covariance() + jnp.outer(self.a_mean, self.b_mean) + + def pearson_correlation(self): + # Note that the 1/(n-1) normalization terms cancel out nicely here. + # TODO: clip? + eps = 0 + # Dead units will have zero variance, which produces NaNs. Convert those to + # zeros with nan_to_num. + return jnp.nan_to_num(self.cov / (jnp.sqrt(self.var_a[:, jnp.newaxis]) + eps) / + (jnp.sqrt(self.var_b) + eps)) + +class OnlineInnerProduct(NamedTuple): + val: jnp.ndarray # (d, d) + + @staticmethod + def init(d: int): + return OnlineInnerProduct(val=jnp.zeros((d, d))) + + def update(self, a_batch, b_batch): + assert a_batch.shape == b_batch.shape + return OnlineInnerProduct(val=self.val + a_batch.T @ b_batch) diff --git a/src/utils.py b/src/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..20c35033376dba39a63309d6b701d6a380cf8d39 --- /dev/null +++ b/src/utils.py @@ -0,0 +1,196 @@ +import operator +import re +import time +from contextlib import contextmanager + +import jax.numpy as jnp +from flax import traverse_util +from flax.core import freeze, unfreeze +from jax import random +from jax.tree_util import tree_reduce, tree_map + +rngmix = lambda rng, x: random.fold_in(rng, hash(x) % 2**16) + +@contextmanager +def timeblock(name): + start = time.time() + try: + yield + finally: + end = time.time() + print(f"{name} took {end - start:.5f} seconds") + +class RngPooper: + """A stateful wrapper around stateless random.PRNGKey's.""" + + def __init__(self, init_rng): + self.rng = init_rng + + def poop(self): + self.rng, rng_key = random.split(self.rng) + return rng_key + +def l1prox(x, alpha): + return jnp.sign(x) * jnp.maximum(0, jnp.abs(x) - alpha) + +def ec2_get_instance_type(): + # See also https://stackoverflow.com/questions/51486405/aws-ec2-command-line-display-instance-type/51486782 + return open("/sys/devices/virtual/dmi/id/product_name").read().strip() + +# Utilities for dealing with flax model parameters +def partition(pred, iterable): + trues = [] + falses = [] + for item in iterable: + if pred(item): + trues.append(item) + else: + falses.append(item) + return trues, falses + +def partition_dict(pred, d): + trues = {} + falses = {} + for k, v in d.items(): + if pred(k): + trues[k] = v + else: + falses[k] = v + return trues, falses + +def flatten_params(params): + return {"/".join(k): v for k, v in traverse_util.flatten_dict(unfreeze(params)).items()} + +def unflatten_params(flat_params): + return freeze( + traverse_util.unflatten_dict({tuple(k.split("/")): v + for k, v in flat_params.items()})) + +def merge_params(a, b): + return unflatten_params({**a, **b}) + +def kmatch(pattern, key): + regex = "^" + i = 0 + while i < len(pattern): + if pattern[i] == "*": + if i + 1 < len(pattern) and pattern[i + 1] == "*": + regex += "(.*)" + i += 2 + else: + regex += "([^\/]*)" + i += 1 + else: + regex += pattern[i] + i += 1 + regex += "$" + return re.fullmatch(regex, key) + +assert kmatch("*", "a") is not None +assert kmatch("*", "a").group(0) == "a" +assert kmatch("*", "a").group(1) == "a" +assert kmatch("abc", "def") is None +assert kmatch("abc/*/ghi", "abc/def/ghi").group(1) == "def" +assert kmatch("abc/**/jkl", "abc/def/ghi/jkl").group(1) == "def/ghi" +assert kmatch("abc/*/jkl", "abc/def/ghi/jkl") is None +assert kmatch("**/*", "abc/def/ghi/jkl").group(1) == "abc/def/ghi" +assert kmatch("**/*", "abc/def/ghi/jkl").group(2) == "jkl" + +def lerp(lam, t1, t2): + return tree_map(lambda a, b: (1 - lam) * a + lam * b, t1, t2) + +def tree_norm(t): + return jnp.sqrt(tree_reduce(operator.add, tree_map(lambda x: jnp.sum(x**2), t))) + +def tree_l2(t1, t2): + return tree_norm(tree_map(lambda x, y: x - y, t1, t2)) + +def slerp(lam, t1, t2): + # See https://en.wikipedia.org/wiki/Slerp + om = jnp.arccos( + tree_reduce(operator.add, tree_map(lambda x, y: jnp.sum(x * y), t1, t2)) / + (tree_norm(t1) * tree_norm(t2))) + sinom = jnp.sin(om) + return tree_map( + lambda x, y: jnp.sin((1 - lam) * om) / sinom * x + jnp.sin(lam * om) / sinom * y, + t1, + t2, + ) + +def compute_weights_cost_matrices(model, params_a, params_b): + """ + Compute weight-based cost matrices for two MoE models based on gating weights and expert parameters. + + Args: + model: ViTModelMoE instance with a get_moe_params method and attributes num_experts, embedding_dim, hidden_dim + params_a: Parameters of the first model + params_b: Parameters of the second model + + Returns: + D: Distance matrix [num_experts, num_experts] based on expert parameters using Frobenius norm + S: Distance matrix [num_experts, num_experts] based on gating weights + """ + moe_params_a = model.get_moe_params(params_a) + moe_params_b = model.get_moe_params(params_b) + num_experts = model.num_experts + + # Gating-based similarity matrix S + gating_kernel_a = moe_params_a['gating_kernel'] + gating_bias_a = moe_params_a['gating_bias'] + gating_kernel_b = moe_params_b['gating_kernel'] + gating_bias_b = moe_params_b['gating_bias'] + + # Center the gating kernels and biases + mean_gating_kernel_a = np.mean(gating_kernel_a, axis=0) + mean_gating_bias_a = np.mean(gating_bias_a) + centered_gating_kernel_a = gating_kernel_a - mean_gating_kernel_a + centered_gating_bias_a = gating_bias_a - mean_gating_bias_a + + mean_gating_kernel_b = np.mean(gating_kernel_b, axis=0) + mean_gating_bias_b = np.mean(gating_bias_b) + centered_gating_kernel_b = gating_kernel_b - mean_gating_kernel_b + centered_gating_bias_b = gating_bias_b - mean_gating_bias_b + + # Construct gating vectors with centered kernels and biases + gating_vectors_a = np.hstack([centered_gating_kernel_a.T, centered_gating_bias_a[:, np.newaxis]]) + gating_vectors_b = np.hstack([centered_gating_kernel_b.T, centered_gating_bias_b[:, np.newaxis]]) + + # Compute the difference between gating vectors + diff_vectors = gating_vectors_a[:, np.newaxis, :] - gating_vectors_b[np.newaxis, :, :] + + # Compute the Euclidean distance + S = np.sqrt(np.sum(diff_vectors ** 2, axis=2)) + + # Expert parameters-based distance matrix D + D = np.zeros((num_experts, num_experts)) + for i in range(num_experts): + W1_a = moe_params_a[f'expert_{i}_layer1_kernel'] + b1_a = moe_params_a[f'expert_{i}_layer1_bias'] + W1p_a = moe_params_a[f'expert_{i}_layer2_kernel'] + b1p_a = moe_params_a[f'expert_{i}_layer2_bias'] + + W_tilde1_a = np.vstack([W1_a, b1_a[np.newaxis, :]]) + W_tilde1p_a = np.vstack([W1p_a, b1p_a[np.newaxis, :]]) + + gram1_a = W_tilde1_a.T @ W_tilde1_a + gram1p_a = W_tilde1p_a @ W_tilde1p_a.T + + for j in range(num_experts): + W1_b = moe_params_b[f'expert_{j}_layer1_kernel'] + b1_b = moe_params_b[f'expert_{j}_layer1_bias'] + W1p_b = moe_params_b[f'expert_{j}_layer2_kernel'] + b1p_b = moe_params_b[f'expert_{j}_layer2_bias'] + + W_tilde1_b = np.vstack([W1_b, b1_b[np.newaxis, :]]) + W_tilde1p_b = np.vstack([W1p_b, b1p_b[np.newaxis, :]]) + + gram1_b = W_tilde1_b.T @ W_tilde1_b + gram1p_b = W_tilde1p_b @ W_tilde1p_b.T + + diff1 = gram1_a - gram1_b + diff1p = gram1p_a - gram1p_b + norm_diff1 = np.linalg.norm(diff1, 'fro') + norm_diff1p = np.linalg.norm(diff1p, 'fro') + D[i, j] = np.sqrt(norm_diff1**2 + norm_diff1p**2) + + return D, S diff --git a/temp/bert.py b/temp/bert.py new file mode 100644 index 0000000000000000000000000000000000000000..7e2760b167b9e22098408d7733e3232d4bd76281 --- /dev/null +++ b/temp/bert.py @@ -0,0 +1,93 @@ +import torch +import numpy as np +import csv +from transformers import BertModel + +# Load model +model = BertModel.from_pretrained("bert-base-uncased") + +hidden_size = model.config.hidden_size +num_heads = model.config.num_attention_heads +head_dim = hidden_size // num_heads + +# Define norms +def compute_norms(A: torch.Tensor): + return { + "1": torch.norm(A, p=1).item(), + "F": torch.norm(A, p="fro").item(), + "*": torch.linalg.svdvals(A).sum().item(), + "2,1": torch.norm(A, dim=0, p=2).sum().item(), + "2,1,T": torch.norm(A.t(), dim=0, p=2).sum().item() + } + +norm_names = ["1","F","*","2,1","2,1,T"] +subcols = ["Q","K","Q/K","V","O","V/O"] + +outfile = "bert_qkvo_norms.csv" + +with open(outfile, "w", newline="") as f: + writer = csv.writer(f) + + for layer_idx, layer in enumerate(model.encoder.layer, start=1): + # Header + header1 = [f"Layer {layer_idx}"] + for n in norm_names: + header1.extend([n,"","","","",""]) + writer.writerow(header1) + + header2 = [""] + for _ in norm_names: + header2.extend(subcols) + writer.writerow(header2) + + rows = [] + + # Extract weights + concat bias as extra column + W_q = torch.cat([layer.attention.self.query.weight.detach(), + layer.attention.self.query.bias.detach().unsqueeze(1)], dim=1) + W_k = torch.cat([layer.attention.self.key.weight.detach(), + layer.attention.self.key.bias.detach().unsqueeze(1)], dim=1) + W_v = torch.cat([layer.attention.self.value.weight.detach(), + layer.attention.self.value.bias.detach().unsqueeze(1)], dim=1) + W_o = torch.cat([layer.attention.output.dense.weight.detach(), + layer.attention.output.dense.bias.detach().unsqueeze(1)], dim=1) + + # Split into heads + W_q_heads = W_q.view(num_heads, head_dim, -1) + W_k_heads = W_k.view(num_heads, head_dim, -1) + W_v_heads = W_v.view(num_heads, head_dim, -1) + W_o_heads = W_o.view(num_heads, head_dim, -1) + + # Head rows + for h in range(num_heads): + row = [f"Head {h+1}"] + for norm in norm_names: + nq = compute_norms(W_q_heads[h])[norm] + nk = compute_norms(W_k_heads[h])[norm] + nv = compute_norms(W_v_heads[h])[norm] + no = compute_norms(W_o_heads[h])[norm] + + qk_ratio = nq/(nk+1e-12) + vo_ratio = nv/(no+1e-12) + + row.extend([ + round(nq,4), + round(nk,4), + round(qk_ratio,4), + round(nv,4), + round(no,4), + round(vo_ratio,4) + ]) + writer.writerow(row) + rows.append(row[1:]) + + # Mean & Std + arr = np.array(rows, dtype=float) + mean = np.round(arr.mean(axis=0),4) + std = np.round(arr.std(axis=0),4) + + writer.writerow(["Mean"] + mean.tolist()) + writer.writerow(["Std"] + std.tolist()) + writer.writerow([]) + +print(f"✅ Saved CSV: {outfile}") diff --git a/temp/gpt2.py b/temp/gpt2.py new file mode 100644 index 0000000000000000000000000000000000000000..b6086dfa7b79d169ce8da0a83f63eb933d560bc3 --- /dev/null +++ b/temp/gpt2.py @@ -0,0 +1,33 @@ +import torch +from transformers import GPT2Model + +# Load pretrained GPT-2 +model = GPT2Model.from_pretrained("openai-community/gpt2") + +ratios = [] + +# Loop through transformer blocks +for i, block in enumerate(model.h): # model.h = list of transformer layers + # GPT2 uses Conv1D for QKV, packed into one matrix + W_qkv = block.attn.c_attn.weight.detach() # shape [768, 2304] + + # Split into Q, K, V + hidden_size = model.config.hidden_size # 768 + W_q, W_k, W_v = W_qkv.split(hidden_size, dim=1) # each [768, 768] + + # Compute norms + l2_q = torch.norm(W_q, p=2) + l2_k = torch.norm(W_k, p=2) + + ratio = l2_q / l2_k + ratios.append(ratio.item()) + + print(f"Layer {i+1}: Q/K L2 ratio = {ratio.item():.4f}") + +# Convert to tensor for stats +ratios_tensor = torch.tensor(ratios) +mean_ratio = torch.mean(ratios_tensor).item() +std_ratio = torch.std(ratios_tensor).item() + +print(f"\nMean Q/K ratio over {len(ratios)} layers: {mean_ratio:.4f}") +print(f"Std Q/K ratio over {len(ratios)} layers: {std_ratio:.4f}") diff --git a/temp/gptj.py b/temp/gptj.py new file mode 100644 index 0000000000000000000000000000000000000000..23e612b1bfd5c035f25e2cd43392cf0892252107 --- /dev/null +++ b/temp/gptj.py @@ -0,0 +1,94 @@ +import torch +import numpy as np +import csv +from transformers import GPTJModel + +# Load GPT-J (weights on CPU with float16) +model = GPTJModel.from_pretrained( + "EleutherAI/gpt-j-6b", + torch_dtype=torch.float16, + low_cpu_mem_usage=True +).cpu() + +hidden_size = model.config.hidden_size +num_heads = model.config.n_head +head_dim = hidden_size // num_heads + +# ---- Norm calculator ---- +def compute_norms(A: torch.Tensor): + A = A.float() # for numerical stability + return { + "1": torch.norm(A, p=1).item(), + "F": torch.norm(A, p="fro").item(), + "*": torch.linalg.svdvals(A).sum().item(), + "2,1": torch.norm(A, dim=0, p=2).sum().item(), + "2,1,T": torch.norm(A.t(), dim=0, p=2).sum().item() + } + +norm_names = ["1","F","*","2,1","2,1,T"] +subcols = ["Q","K","Q/K","V","O","V/O"] + +outfile = "gptj_qkvo_norms.csv" + +with open(outfile, "w", newline="") as f: + writer = csv.writer(f) + + for layer_idx, layer in enumerate(model.h, start=1): + # ---- Headers ---- + header1 = [f"Layer {layer_idx}"] + for n in norm_names: + header1.extend([n,"","","","",""]) + writer.writerow(header1) + + header2 = [""] + for _ in norm_names: + header2.extend(subcols) + writer.writerow(header2) + + rows = [] + + # ---- Extract weights (no bias in GPT-J) ---- + W_q = layer.attn.q_proj.weight.detach().cpu() + W_k = layer.attn.k_proj.weight.detach().cpu() + W_v = layer.attn.v_proj.weight.detach().cpu() + W_o = layer.attn.out_proj.weight.detach().cpu() + + # ---- Split into heads ---- + W_q_heads = W_q.view(num_heads, head_dim, -1) + W_k_heads = W_k.view(num_heads, head_dim, -1) + W_v_heads = W_v.view(num_heads, head_dim, -1) + W_o_heads = W_o.view(num_heads, head_dim, -1) + + # ---- Per-head ---- + for h in range(num_heads): + row = [f"Head {h+1}"] + for norm in norm_names: + nq = compute_norms(W_q_heads[h])[norm] + nk = compute_norms(W_k_heads[h])[norm] + nv = compute_norms(W_v_heads[h])[norm] + no = compute_norms(W_o_heads[h])[norm] + + qk_ratio = nq / (nk + 1e-12) + vo_ratio = nv / (no + 1e-12) + + row.extend([ + round(nq,4), + round(nk,4), + round(qk_ratio,4), + round(nv,4), + round(no,4), + round(vo_ratio,4) + ]) + writer.writerow(row) + rows.append(row[1:]) + + # ---- Mean & Std ---- + arr = np.array(rows, dtype=float) + mean = np.round(arr.mean(axis=0),4) + std = np.round(arr.std(axis=0),4) + + writer.writerow(["Mean"] + mean.tolist()) + writer.writerow(["Std"] + std.tolist()) + writer.writerow([]) + +print(f"✅ Saved CSV: {outfile}") diff --git a/temp/qwen.py b/temp/qwen.py new file mode 100644 index 0000000000000000000000000000000000000000..639f6a8c406ad412ba4255b8c3ec9669bff488f6 --- /dev/null +++ b/temp/qwen.py @@ -0,0 +1,109 @@ +import torch +from transformers import AutoTokenizer, AutoModelForCausalLM +import csv +import numpy as np + +# ------------------ Load model ------------------ +model_name = "Qwen/Qwen2-0.5B" +model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32) + +# ------------------ Detect architecture -------- + +print(model) +# Compute norms +def compute_norms(A: torch.Tensor): + print(A.shape) + norms = {} + norms["1"] = torch.norm(A, p=1).item() + norms["inf"] = torch.norm(A, p=float("inf")).item() + norms["2"] = torch.linalg.svdvals(A).max().item() + norms["F"] = torch.norm(A, p="fro").item() + norms["*"] = torch.linalg.svdvals(A).sum().item() + norms["max"] = A.abs().max().item() + norms["2,1"] = torch.norm(A, dim=0, p=2).sum().item() + norms["2,1,T"] = torch.norm(A.t(), dim=0, p=2).sum().item() + return norms + +def dict_of_lists(norm_list): + result = {k: [] for k in norm_list[0].keys()} + for d in norm_list: + for k, v in d.items(): + result[k].append(v) + return result + +# Store norms and ratios +all_norms_Q, all_norms_K, norm_ratios = [], [], [] + +for layer in model.model.layers: + W_q = layer.self_attn.q_proj.weight.detach() + W_k = layer.self_attn.k_proj.weight.detach() + if W_k.shape[0] != W_q.shape[0]: + repeat_factor = W_q.shape[0] // W_k.shape[0] + W_k = W_k.repeat(repeat_factor, 1) + norms_q = compute_norms(W_q) + norms_k = compute_norms(W_k) + + all_norms_Q.append(norms_q) + all_norms_K.append(norms_k) + + ratio_dict = {k: norms_q[k] / norms_k[k] for k in norms_q.keys()} + norm_ratios.append(ratio_dict) + +Q_dict = dict_of_lists(all_norms_Q) +K_dict = dict_of_lists(all_norms_K) +R_dict = dict_of_lists(norm_ratios) + +sections = [ + (["1", "inf", "2", "F"], ["1", "inf", "2", "F"]), + (["*", "max", "2,1", "2,1,T"], ["*", "max", "2,1", "2,1,T"]) +] + +with open("qwen2.5_norms.csv", "w", newline="") as f: + writer = csv.writer(f) + + for idx, (main_norms, keys) in enumerate(sections): + # Header 1 + header1 = ["Layer"] + for norm in main_norms: + header1.extend([norm, "", ""]) + writer.writerow(header1) + + # Header 2 + header2 = [""] + for _ in main_norms: + header2.extend(["Q", "K", "Q/K"]) + writer.writerow(header2) + + # Per-layer rows + for i in range(len(all_norms_Q)): + row = [i + 1] + for key in keys: + row.extend([ + round(Q_dict[key][i], 4), + round(K_dict[key][i], 4), + round(R_dict[key][i], 4) + ]) + writer.writerow(row) + + # Mean and Std + mean_row = ["Mean and Std"] + std_row = [""] + + for key in keys: + q_vals = np.array(Q_dict[key]) + k_vals = np.array(K_dict[key]) + r_vals = np.array(R_dict[key]) + + mean_row.extend([ + round(q_vals.mean(), 4), + round(k_vals.mean(), 4), + round(r_vals.mean(), 4) + ]) + std_row.extend([ + round(q_vals.std(), 4), + round(k_vals.std(), 4), + round(r_vals.std(), 4) + ]) + + writer.writerow(mean_row) + writer.writerow(std_row) diff --git a/temp/roformer.py b/temp/roformer.py new file mode 100644 index 0000000000000000000000000000000000000000..2cc221cca2ce468705787a9ea97cf344979a51c3 --- /dev/null +++ b/temp/roformer.py @@ -0,0 +1,93 @@ +import torch +import numpy as np +import csv +from transformers import RoFormerModel + +# Load RoFormer +model = RoFormerModel.from_pretrained("junnyu/roformer_chinese_base") + +hidden_size = model.config.hidden_size +num_heads = model.config.num_attention_heads +head_dim = hidden_size // num_heads + +# Norm calculator +def compute_norms(A: torch.Tensor): + return { + "1": torch.norm(A, p=1).item(), + "F": torch.norm(A, p="fro").item(), + "*": torch.linalg.svdvals(A).sum().item(), + "2,1": torch.norm(A, dim=0, p=2).sum().item(), + "2,1,T": torch.norm(A.t(), dim=0, p=2).sum().item() + } + +norm_names = ["1","F","*","2,1","2,1,T"] +subcols = ["Q","K","Q/K","V","O","V/O"] + +outfile = "roformer_qkvo_norms.csv" + +with open(outfile, "w", newline="") as f: + writer = csv.writer(f) + + for layer_idx, layer in enumerate(model.encoder.layer, start=1): + # Header rows + header1 = [f"Layer {layer_idx}"] + for n in norm_names: + header1.extend([n,"","","","",""]) + writer.writerow(header1) + + header2 = [""] + for _ in norm_names: + header2.extend(subcols) + writer.writerow(header2) + + rows = [] + + # ---- Extract weights + biases ---- + W_q = torch.cat([layer.attention.self.query.weight.detach(), + layer.attention.self.query.bias.detach().unsqueeze(1)], dim=1) + W_k = torch.cat([layer.attention.self.key.weight.detach(), + layer.attention.self.key.bias.detach().unsqueeze(1)], dim=1) + W_v = torch.cat([layer.attention.self.value.weight.detach(), + layer.attention.self.value.bias.detach().unsqueeze(1)], dim=1) + W_o = torch.cat([layer.attention.output.dense.weight.detach(), + layer.attention.output.dense.bias.detach().unsqueeze(1)], dim=1) + + # ---- Split into heads ---- + W_q_heads = W_q.view(num_heads, head_dim, -1) + W_k_heads = W_k.view(num_heads, head_dim, -1) + W_v_heads = W_v.view(num_heads, head_dim, -1) + W_o_heads = W_o.view(num_heads, head_dim, -1) + + # ---- Per-head norms ---- + for h in range(num_heads): + row = [f"Head {h+1}"] + for norm in norm_names: + nq = compute_norms(W_q_heads[h])[norm] + nk = compute_norms(W_k_heads[h])[norm] + nv = compute_norms(W_v_heads[h])[norm] + no = compute_norms(W_o_heads[h])[norm] + + qk_ratio = nq/(nk+1e-12) + vo_ratio = nv/(no+1e-12) + + row.extend([ + round(nq,4), + round(nk,4), + round(qk_ratio,4), + round(nv,4), + round(no,4), + round(vo_ratio,4) + ]) + writer.writerow(row) + rows.append(row[1:]) + + # ---- Mean & Std ---- + arr = np.array(rows, dtype=float) + mean = np.round(arr.mean(axis=0),4) + std = np.round(arr.std(axis=0),4) + + writer.writerow(["Mean"] + mean.tolist()) + writer.writerow(["Std"] + std.tolist()) + writer.writerow([]) + +print(f"✅ Saved CSV: {outfile}") diff --git a/temp/vit.py b/temp/vit.py new file mode 100644 index 0000000000000000000000000000000000000000..4b74cd4af5f0f36e415ce2dcd1614397c71e4a72 --- /dev/null +++ b/temp/vit.py @@ -0,0 +1,92 @@ +import torch +import numpy as np +import csv +from transformers import ViTModel + +# Load ViT +model = ViTModel.from_pretrained("google/vit-base-patch16-224") + +hidden_size = model.config.hidden_size +num_heads = model.config.num_attention_heads +head_dim = hidden_size // num_heads + +# Norm calculator +def compute_norms(A: torch.Tensor): + return { + "1": torch.norm(A, p=1).item(), + "F": torch.norm(A, p="fro").item(), + "*": torch.linalg.svdvals(A).sum().item(), + "2,1": torch.norm(A, dim=0, p=2).sum().item(), + "2,1,T": torch.norm(A.t(), dim=0, p=2).sum().item() + } + +norm_names = ["1","F","*","2,1","2,1,T"] +subcols = ["Q","K","Q/K","V","O","V/O"] + +outfile = "vit_qkvo_norms.csv" + +with open(outfile, "w", newline="") as f: + writer = csv.writer(f) + + for layer_idx, layer in enumerate(model.encoder.layer, start=1): + # Header rows + header1 = [f"Layer {layer_idx}"] + for n in norm_names: + header1.extend([n,"","","","",""]) + writer.writerow(header1) + + header2 = [""] + for _ in norm_names: + header2.extend(subcols) + writer.writerow(header2) + + rows = [] + + # Extract weights + biases + W_q = torch.cat([layer.attention.attention.query.weight.detach(), + layer.attention.attention.query.bias.detach().unsqueeze(1)], dim=1) + W_k = torch.cat([layer.attention.attention.key.weight.detach(), + layer.attention.attention.key.bias.detach().unsqueeze(1)], dim=1) + W_v = torch.cat([layer.attention.attention.value.weight.detach(), + layer.attention.attention.value.bias.detach().unsqueeze(1)], dim=1) + W_o = torch.cat([layer.attention.output.dense.weight.detach(), + layer.attention.output.dense.bias.detach().unsqueeze(1)], dim=1) + + # Reshape into heads + W_q_heads = W_q.view(num_heads, head_dim, -1) + W_k_heads = W_k.view(num_heads, head_dim, -1) + W_v_heads = W_v.view(num_heads, head_dim, -1) + W_o_heads = W_o.view(num_heads, head_dim, -1) + + for h in range(num_heads): + row = [f"Head {h+1}"] + for norm in norm_names: + nq = compute_norms(W_q_heads[h])[norm] + nk = compute_norms(W_k_heads[h])[norm] + nv = compute_norms(W_v_heads[h])[norm] + no = compute_norms(W_o_heads[h])[norm] + + qk_ratio = nq/(nk+1e-12) + vo_ratio = nv/(no+1e-12) + + row.extend([ + round(nq,4), + round(nk,4), + round(qk_ratio,4), + round(nv,4), + round(no,4), + round(vo_ratio,4) + ]) + writer.writerow(row) + rows.append(row[1:]) + + # Mean & Std rows + arr = np.array(rows, dtype=float) + mean = np.round(arr.mean(axis=0),4) + std = np.round(arr.std(axis=0),4) + + writer.writerow(["Mean"] + mean.tolist()) + writer.writerow(["Std"] + std.tolist()) + writer.writerow([]) + +print(f"✅ Saved CSV: {outfile}") diff --git a/wandb/debug-internal.log b/wandb/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..0658c12996a5f0c436205033c2de8f24e45ef307 --- /dev/null +++ b/wandb/debug-internal.log @@ -0,0 +1,7 @@ +{"time":"2025-12-26T09:41:08.773359688+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-26T09:41:09.249853643+07:00","level":"INFO","msg":"stream: created new stream","id":"svj9s262"} +{"time":"2025-12-26T09:41:09.249933283+07:00","level":"INFO","msg":"handler: started","stream_id":"svj9s262"} +{"time":"2025-12-26T09:41:09.250011964+07:00","level":"INFO","msg":"stream: started","id":"svj9s262"} +{"time":"2025-12-26T09:41:09.250090029+07:00","level":"INFO","msg":"writer: started","stream_id":"svj9s262"} +{"time":"2025-12-26T09:41:09.250121649+07:00","level":"INFO","msg":"sender: started","stream_id":"svj9s262"} +{"time":"2025-12-26T09:41:29.560166541+07:00","level":"INFO","msg":"stream: closing","id":"svj9s262"} diff --git a/wandb/debug.log b/wandb/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..1b4ceac7a0a901622a983c980755b3d7002fd51d --- /dev/null +++ b/wandb/debug.log @@ -0,0 +1,348 @@ +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_setup.py:_flush():80] Configure stats pid to 1318853 +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251226_094108-svj9s262/logs/debug.log +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251226_094108-svj9s262/logs/debug-internal.log +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_init.py:init():840] calling init triggers +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_init.py:init():888] starting backend +2025-12-26 09:41:08,759 INFO MainThread:1318853 [wandb_init.py:init():891] sending inform_init request +2025-12-26 09:41:08,766 INFO MainThread:1318853 [wandb_init.py:init():899] backend started and connected +2025-12-26 09:41:08,768 INFO MainThread:1318853 [wandb_init.py:init():969] updated telemetry +2025-12-26 09:41:08,772 INFO MainThread:1318853 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-26 09:41:09,748 INFO MainThread:1318853 [wandb_init.py:init():1040] starting run threads in backend +2025-12-26 09:41:09,895 INFO MainThread:1318853 [wandb_run.py:_console_start():2504] atexit reg +2025-12-26 09:41:09,895 INFO MainThread:1318853 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-26 09:41:09,895 INFO MainThread:1318853 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-26 09:41:09,895 INFO MainThread:1318853 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-26 09:41:09,898 INFO MainThread:1318853 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-26 09:41:29,560 INFO wandb-AsyncioManager-main:1318853 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-26 09:41:29,560 INFO wandb-AsyncioManager-main:1318853 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. +2025-12-26 09:41:29,931 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:29,938 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:29,939 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:29,944 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:29,975 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,011 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,022 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,091 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,108 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,108 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,110 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,112 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,112 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,112 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,117 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,118 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,118 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,119 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,119 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost diff --git a/wandb/offline-run-20251225_114356-lqhojgc3/files/code/src/imagenet/train_model.py b/wandb/offline-run-20251225_114356-lqhojgc3/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..ad886085b66d9b18d44a70fa31214a8ca254d6ee --- /dev/null +++ b/wandb/offline-run-20251225_114356-lqhojgc3/files/code/src/imagenet/train_model.py @@ -0,0 +1,262 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/offline-run-20251225_114356-lqhojgc3/files/diff.patch b/wandb/offline-run-20251225_114356-lqhojgc3/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..64066b87ba66f006923a1de989a5c15561cbb8ff --- /dev/null +++ b/wandb/offline-run-20251225_114356-lqhojgc3/files/diff.patch @@ -0,0 +1,134 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..b4f18b4 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..f57f9d2 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': + dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..ad88608 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,13 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +212,7 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/offline-run-20251225_114356-lqhojgc3/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/offline-run-20251225_114356-lqhojgc3/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..64066b87ba66f006923a1de989a5c15561cbb8ff --- /dev/null +++ b/wandb/offline-run-20251225_114356-lqhojgc3/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,134 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..b4f18b4 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..f57f9d2 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': + dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..ad88608 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,13 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +212,7 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/offline-run-20251225_114356-lqhojgc3/files/requirements.txt b/wandb/offline-run-20251225_114356-lqhojgc3/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/offline-run-20251225_114356-lqhojgc3/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/offline-run-20251225_114356-lqhojgc3/logs/debug-core.log b/wandb/offline-run-20251225_114356-lqhojgc3/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..0262873a0910980adbdb8c999c99020a645abbe8 --- /dev/null +++ b/wandb/offline-run-20251225_114356-lqhojgc3/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-25T11:43:56.217537934+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmpyk7_pgac/port-24560.txt","pid":24560,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T11:43:56.218519377+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":24560} +{"time":"2025-12-25T11:43:56.218516013+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-24560-24683-1361265209/socket","Net":"unix"}} +{"time":"2025-12-25T11:43:56.398573871+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T11:43:56.405592008+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"lqhojgc3","id":"1(@)"} +{"time":"2025-12-25T11:43:56.638231109+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"lqhojgc3","id":"1(@)"} +{"time":"2025-12-25T11:43:57.281184892+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T11:43:57.281261504+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T11:43:57.281267154+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T11:43:57.281308824+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-24560-24683-1361265209/socket","Net":"unix"}} +{"time":"2025-12-25T11:43:57.281321926+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T11:43:57.285169554+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-25T11:43:57.285181755+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-25T11:43:57.285184768+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/offline-run-20251225_114356-lqhojgc3/logs/debug-internal.log b/wandb/offline-run-20251225_114356-lqhojgc3/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..c5f849ffc4a4c15b3bf3f3327ee5a3e94e80c461 --- /dev/null +++ b/wandb/offline-run-20251225_114356-lqhojgc3/logs/debug-internal.log @@ -0,0 +1,12 @@ +{"time":"2025-12-25T11:43:56.409073241+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T11:43:56.638007272+07:00","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"} +{"time":"2025-12-25T11:43:56.638099933+07:00","level":"INFO","msg":"stream: created new stream","id":"lqhojgc3"} +{"time":"2025-12-25T11:43:56.638146419+07:00","level":"INFO","msg":"handler: started","stream_id":"lqhojgc3"} +{"time":"2025-12-25T11:43:56.638224975+07:00","level":"INFO","msg":"stream: started","id":"lqhojgc3"} +{"time":"2025-12-25T11:43:56.638255966+07:00","level":"INFO","msg":"writer: started","stream_id":"lqhojgc3"} +{"time":"2025-12-25T11:43:56.63829098+07:00","level":"INFO","msg":"sender: started","stream_id":"lqhojgc3"} +{"time":"2025-12-25T11:43:56.638876417+07:00","level":"WARN","msg":"runupserter: server does not expand metric globs but the x_server_side_expand_glob_metrics setting is set; ignoring"} +{"time":"2025-12-25T11:43:57.281212192+07:00","level":"INFO","msg":"stream: closing","id":"lqhojgc3"} +{"time":"2025-12-25T11:43:57.281398772+07:00","level":"INFO","msg":"handler: closed","stream_id":"lqhojgc3"} +{"time":"2025-12-25T11:43:57.281437489+07:00","level":"INFO","msg":"sender: closed","stream_id":"lqhojgc3"} +{"time":"2025-12-25T11:43:57.281441129+07:00","level":"INFO","msg":"stream: closed","id":"lqhojgc3"} diff --git a/wandb/offline-run-20251225_114356-lqhojgc3/logs/debug.log b/wandb/offline-run-20251225_114356-lqhojgc3/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..6f082b05b42128fbb5a57f62ca157fe2b6eec51b --- /dev/null +++ b/wandb/offline-run-20251225_114356-lqhojgc3/logs/debug.log @@ -0,0 +1,23 @@ +2025-12-25 11:43:56,121 INFO MainThread:24560 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 11:43:56,121 INFO MainThread:24560 [wandb_setup.py:_flush():80] Configure stats pid to 24560 +2025-12-25 11:43:56,121 INFO MainThread:24560 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 11:43:56,121 INFO MainThread:24560 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 11:43:56,121 INFO MainThread:24560 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 11:43:56,122 INFO MainThread:24560 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/offline-run-20251225_114356-lqhojgc3/logs/debug.log +2025-12-25 11:43:56,122 INFO MainThread:24560 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/offline-run-20251225_114356-lqhojgc3/logs/debug-internal.log +2025-12-25 11:43:56,122 INFO MainThread:24560 [wandb_init.py:init():840] calling init triggers +2025-12-25 11:43:56,122 INFO MainThread:24560 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 11:43:56,122 INFO MainThread:24560 [wandb_init.py:init():888] starting backend +2025-12-25 11:43:56,398 INFO MainThread:24560 [wandb_init.py:init():891] sending inform_init request +2025-12-25 11:43:56,403 INFO MainThread:24560 [wandb_init.py:init():899] backend started and connected +2025-12-25 11:43:56,404 INFO MainThread:24560 [wandb_init.py:init():969] updated telemetry +2025-12-25 11:43:56,409 INFO MainThread:24560 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 11:43:56,640 INFO MainThread:24560 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 11:43:56,726 INFO MainThread:24560 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 11:43:56,727 INFO MainThread:24560 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 11:43:56,727 INFO MainThread:24560 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 11:43:56,727 INFO MainThread:24560 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 11:43:56,728 INFO MainThread:24560 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-25 11:43:57,281 INFO wandb-AsyncioManager-main:24560 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-25 11:43:57,281 INFO wandb-AsyncioManager-main:24560 [mailbox.py:close():137] Closing mailbox, abandoning 0 handles. diff --git a/wandb/offline-run-20251225_114356-lqhojgc3/run-lqhojgc3.wandb b/wandb/offline-run-20251225_114356-lqhojgc3/run-lqhojgc3.wandb new file mode 100644 index 0000000000000000000000000000000000000000..1e980762f445dd3bb74b0c77636288778b95b161 Binary files /dev/null and b/wandb/offline-run-20251225_114356-lqhojgc3/run-lqhojgc3.wandb differ diff --git a/wandb/offline-run-20251225_114434-s87ka52z/files/code/src/imagenet/train_model.py b/wandb/offline-run-20251225_114434-s87ka52z/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..ad886085b66d9b18d44a70fa31214a8ca254d6ee --- /dev/null +++ b/wandb/offline-run-20251225_114434-s87ka52z/files/code/src/imagenet/train_model.py @@ -0,0 +1,262 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/offline-run-20251225_114434-s87ka52z/files/diff.patch b/wandb/offline-run-20251225_114434-s87ka52z/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..64066b87ba66f006923a1de989a5c15561cbb8ff --- /dev/null +++ b/wandb/offline-run-20251225_114434-s87ka52z/files/diff.patch @@ -0,0 +1,134 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..b4f18b4 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..f57f9d2 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': + dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..ad88608 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,13 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +212,7 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/offline-run-20251225_114434-s87ka52z/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/offline-run-20251225_114434-s87ka52z/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..64066b87ba66f006923a1de989a5c15561cbb8ff --- /dev/null +++ b/wandb/offline-run-20251225_114434-s87ka52z/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,134 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..b4f18b4 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..f57f9d2 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': + dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..ad88608 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,13 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +212,7 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/offline-run-20251225_114434-s87ka52z/files/requirements.txt b/wandb/offline-run-20251225_114434-s87ka52z/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/offline-run-20251225_114434-s87ka52z/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/offline-run-20251225_114434-s87ka52z/logs/debug-core.log b/wandb/offline-run-20251225_114434-s87ka52z/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..26a5fab60a6ecf60795303ad55d32ff0c3fca6aa --- /dev/null +++ b/wandb/offline-run-20251225_114434-s87ka52z/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-25T11:44:34.406154776+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmp088o7kqg/port-25131.txt","pid":25131,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T11:44:34.406712586+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-25131-25228-1892185799/socket","Net":"unix"}} +{"time":"2025-12-25T11:44:34.406799442+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":25131} +{"time":"2025-12-25T11:44:34.601505879+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T11:44:34.609396349+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"s87ka52z","id":"1(@)"} +{"time":"2025-12-25T11:44:34.855842198+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"s87ka52z","id":"1(@)"} +{"time":"2025-12-25T11:44:35.514063336+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T11:44:35.514103619+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T11:44:35.514146022+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T11:44:35.514158034+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T11:44:35.514229205+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-25131-25228-1892185799/socket","Net":"unix"}} +{"time":"2025-12-25T11:44:35.518860522+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-25T11:44:35.518879774+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-25T11:44:35.518883277+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/offline-run-20251225_114434-s87ka52z/logs/debug-internal.log b/wandb/offline-run-20251225_114434-s87ka52z/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..4602714e8e8a7c8f23516e9d4b411ab38802dcce --- /dev/null +++ b/wandb/offline-run-20251225_114434-s87ka52z/logs/debug-internal.log @@ -0,0 +1,12 @@ +{"time":"2025-12-25T11:44:34.612668483+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T11:44:34.855499444+07:00","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"} +{"time":"2025-12-25T11:44:34.855636737+07:00","level":"INFO","msg":"stream: created new stream","id":"s87ka52z"} +{"time":"2025-12-25T11:44:34.855700942+07:00","level":"INFO","msg":"handler: started","stream_id":"s87ka52z"} +{"time":"2025-12-25T11:44:34.855832563+07:00","level":"INFO","msg":"stream: started","id":"s87ka52z"} +{"time":"2025-12-25T11:44:34.855875788+07:00","level":"INFO","msg":"writer: started","stream_id":"s87ka52z"} +{"time":"2025-12-25T11:44:34.855895152+07:00","level":"INFO","msg":"sender: started","stream_id":"s87ka52z"} +{"time":"2025-12-25T11:44:34.856589172+07:00","level":"WARN","msg":"runupserter: server does not expand metric globs but the x_server_side_expand_glob_metrics setting is set; ignoring"} +{"time":"2025-12-25T11:44:35.514098501+07:00","level":"INFO","msg":"stream: closing","id":"s87ka52z"} +{"time":"2025-12-25T11:44:35.514231609+07:00","level":"INFO","msg":"handler: closed","stream_id":"s87ka52z"} +{"time":"2025-12-25T11:44:35.514266476+07:00","level":"INFO","msg":"sender: closed","stream_id":"s87ka52z"} +{"time":"2025-12-25T11:44:35.514270407+07:00","level":"INFO","msg":"stream: closed","id":"s87ka52z"} diff --git a/wandb/offline-run-20251225_114434-s87ka52z/logs/debug.log b/wandb/offline-run-20251225_114434-s87ka52z/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..adc578275051e44f984a9bfe5d14fbabf8ad2ab1 --- /dev/null +++ b/wandb/offline-run-20251225_114434-s87ka52z/logs/debug.log @@ -0,0 +1,23 @@ +2025-12-25 11:44:34,353 INFO MainThread:25131 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 11:44:34,353 INFO MainThread:25131 [wandb_setup.py:_flush():80] Configure stats pid to 25131 +2025-12-25 11:44:34,354 INFO MainThread:25131 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 11:44:34,354 INFO MainThread:25131 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 11:44:34,354 INFO MainThread:25131 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 11:44:34,354 INFO MainThread:25131 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/offline-run-20251225_114434-s87ka52z/logs/debug.log +2025-12-25 11:44:34,354 INFO MainThread:25131 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/offline-run-20251225_114434-s87ka52z/logs/debug-internal.log +2025-12-25 11:44:34,354 INFO MainThread:25131 [wandb_init.py:init():840] calling init triggers +2025-12-25 11:44:34,354 INFO MainThread:25131 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 11:44:34,354 INFO MainThread:25131 [wandb_init.py:init():888] starting backend +2025-12-25 11:44:34,601 INFO MainThread:25131 [wandb_init.py:init():891] sending inform_init request +2025-12-25 11:44:34,606 INFO MainThread:25131 [wandb_init.py:init():899] backend started and connected +2025-12-25 11:44:34,608 INFO MainThread:25131 [wandb_init.py:init():969] updated telemetry +2025-12-25 11:44:34,614 INFO MainThread:25131 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 11:44:34,858 INFO MainThread:25131 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 11:44:34,964 INFO MainThread:25131 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 11:44:34,964 INFO MainThread:25131 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 11:44:34,964 INFO MainThread:25131 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 11:44:34,964 INFO MainThread:25131 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 11:44:34,966 INFO MainThread:25131 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-25 11:44:35,514 INFO wandb-AsyncioManager-main:25131 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-25 11:44:35,514 INFO wandb-AsyncioManager-main:25131 [mailbox.py:close():137] Closing mailbox, abandoning 0 handles. diff --git a/wandb/offline-run-20251225_114434-s87ka52z/run-s87ka52z.wandb b/wandb/offline-run-20251225_114434-s87ka52z/run-s87ka52z.wandb new file mode 100644 index 0000000000000000000000000000000000000000..4af47fceb07af6bc9e93d7a6646a5d1e8993f258 Binary files /dev/null and b/wandb/offline-run-20251225_114434-s87ka52z/run-s87ka52z.wandb differ diff --git a/wandb/offline-run-20251225_114458-oskki2mm/files/code/src/imagenet/train_model.py b/wandb/offline-run-20251225_114458-oskki2mm/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..ad886085b66d9b18d44a70fa31214a8ca254d6ee --- /dev/null +++ b/wandb/offline-run-20251225_114458-oskki2mm/files/code/src/imagenet/train_model.py @@ -0,0 +1,262 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/offline-run-20251225_114458-oskki2mm/files/diff.patch b/wandb/offline-run-20251225_114458-oskki2mm/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..0949c1f0dce1fb3d28b8050d5983a6a24963cd62 --- /dev/null +++ b/wandb/offline-run-20251225_114458-oskki2mm/files/diff.patch @@ -0,0 +1,135 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..ad88608 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,13 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +212,7 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/offline-run-20251225_114458-oskki2mm/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/offline-run-20251225_114458-oskki2mm/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..0949c1f0dce1fb3d28b8050d5983a6a24963cd62 --- /dev/null +++ b/wandb/offline-run-20251225_114458-oskki2mm/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,135 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..ad88608 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,13 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +212,7 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/offline-run-20251225_114458-oskki2mm/files/requirements.txt b/wandb/offline-run-20251225_114458-oskki2mm/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/offline-run-20251225_114458-oskki2mm/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/offline-run-20251225_114458-oskki2mm/logs/debug-core.log b/wandb/offline-run-20251225_114458-oskki2mm/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..e4b8017d1834fc2bb370a4b5f0763241b2fb1670 --- /dev/null +++ b/wandb/offline-run-20251225_114458-oskki2mm/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-25T11:44:58.739659184+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmppmsj04re/port-25584.txt","pid":25584,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T11:44:58.74037446+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":25584} +{"time":"2025-12-25T11:44:58.740371729+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-25584-25698-3802468323/socket","Net":"unix"}} +{"time":"2025-12-25T11:44:58.934183051+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T11:44:58.93868481+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"oskki2mm","id":"1(@)"} +{"time":"2025-12-25T11:44:59.070325632+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"oskki2mm","id":"1(@)"} +{"time":"2025-12-25T11:50:47.615177965+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T11:50:47.615370049+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T11:50:47.61544202+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T11:50:47.615456719+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T11:50:47.615619438+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-25584-25698-3802468323/socket","Net":"unix"}} +{"time":"2025-12-25T11:50:47.620735905+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-25T11:50:47.620932073+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-25T11:50:47.620955512+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/offline-run-20251225_114458-oskki2mm/logs/debug-internal.log b/wandb/offline-run-20251225_114458-oskki2mm/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..ced86eb4dee816761bc985398919cff877ba861d --- /dev/null +++ b/wandb/offline-run-20251225_114458-oskki2mm/logs/debug-internal.log @@ -0,0 +1,12 @@ +{"time":"2025-12-25T11:44:58.940857464+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T11:44:59.070083869+07:00","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"} +{"time":"2025-12-25T11:44:59.07016152+07:00","level":"INFO","msg":"stream: created new stream","id":"oskki2mm"} +{"time":"2025-12-25T11:44:59.070211212+07:00","level":"INFO","msg":"handler: started","stream_id":"oskki2mm"} +{"time":"2025-12-25T11:44:59.070319928+07:00","level":"INFO","msg":"stream: started","id":"oskki2mm"} +{"time":"2025-12-25T11:44:59.070351002+07:00","level":"INFO","msg":"writer: started","stream_id":"oskki2mm"} +{"time":"2025-12-25T11:44:59.070367252+07:00","level":"INFO","msg":"sender: started","stream_id":"oskki2mm"} +{"time":"2025-12-25T11:44:59.070859408+07:00","level":"WARN","msg":"runupserter: server does not expand metric globs but the x_server_side_expand_glob_metrics setting is set; ignoring"} +{"time":"2025-12-25T11:50:47.615333425+07:00","level":"INFO","msg":"stream: closing","id":"oskki2mm"} +{"time":"2025-12-25T11:50:47.615778033+07:00","level":"INFO","msg":"handler: closed","stream_id":"oskki2mm"} +{"time":"2025-12-25T11:50:47.615840971+07:00","level":"INFO","msg":"sender: closed","stream_id":"oskki2mm"} +{"time":"2025-12-25T11:50:47.615850521+07:00","level":"INFO","msg":"stream: closed","id":"oskki2mm"} diff --git a/wandb/offline-run-20251225_114458-oskki2mm/logs/debug.log b/wandb/offline-run-20251225_114458-oskki2mm/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..4f56f5541f7cd89153d7cffa72ed601693f1608d --- /dev/null +++ b/wandb/offline-run-20251225_114458-oskki2mm/logs/debug.log @@ -0,0 +1,36 @@ +2025-12-25 11:44:58,689 INFO MainThread:25584 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 11:44:58,689 INFO MainThread:25584 [wandb_setup.py:_flush():80] Configure stats pid to 25584 +2025-12-25 11:44:58,689 INFO MainThread:25584 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 11:44:58,689 INFO MainThread:25584 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 11:44:58,689 INFO MainThread:25584 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 11:44:58,689 INFO MainThread:25584 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/offline-run-20251225_114458-oskki2mm/logs/debug.log +2025-12-25 11:44:58,689 INFO MainThread:25584 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/offline-run-20251225_114458-oskki2mm/logs/debug-internal.log +2025-12-25 11:44:58,689 INFO MainThread:25584 [wandb_init.py:init():840] calling init triggers +2025-12-25 11:44:58,689 INFO MainThread:25584 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 11:44:58,689 INFO MainThread:25584 [wandb_init.py:init():888] starting backend +2025-12-25 11:44:58,934 INFO MainThread:25584 [wandb_init.py:init():891] sending inform_init request +2025-12-25 11:44:58,936 INFO MainThread:25584 [wandb_init.py:init():899] backend started and connected +2025-12-25 11:44:58,937 INFO MainThread:25584 [wandb_init.py:init():969] updated telemetry +2025-12-25 11:44:58,940 INFO MainThread:25584 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 11:44:59,072 INFO MainThread:25584 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 11:44:59,161 INFO MainThread:25584 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 11:44:59,161 INFO MainThread:25584 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 11:44:59,161 INFO MainThread:25584 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 11:44:59,161 INFO MainThread:25584 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 11:44:59,162 INFO MainThread:25584 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-25 11:50:47,615 INFO wandb-AsyncioManager-main:25584 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-25 11:50:47,615 INFO wandb-AsyncioManager-main:25584 [mailbox.py:close():137] Closing mailbox, abandoning 0 handles. +2025-12-25 11:50:47,912 ERROR wandb-AsyncioManager-main:25584 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost diff --git a/wandb/offline-run-20251225_114458-oskki2mm/run-oskki2mm.wandb b/wandb/offline-run-20251225_114458-oskki2mm/run-oskki2mm.wandb new file mode 100644 index 0000000000000000000000000000000000000000..7ab6fa567f2972625801186a278c9402f58977b2 --- /dev/null +++ b/wandb/offline-run-20251225_114458-oskki2mm/run-oskki2mm.wandb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:23d5df638fa23229fa9c2e000aa3ab782083b9b1846fdb19d2c55d4c94febb4a +size 276933 diff --git a/wandb/run-20251225_113509-605ahtng/files/wandb-summary.json b/wandb/run-20251225_113509-605ahtng/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..b0a620d0c1047a4dd8a400939b6da246ed8063a7 --- /dev/null +++ b/wandb/run-20251225_113509-605ahtng/files/wandb-summary.json @@ -0,0 +1 @@ +{"_wandb":{"runtime":0},"_runtime":0} \ No newline at end of file diff --git a/wandb/run-20251225_113509-605ahtng/logs/debug-core.log b/wandb/run-20251225_113509-605ahtng/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..3e3d7550f3268e20da8e1e415522a1716382d5b6 --- /dev/null +++ b/wandb/run-20251225_113509-605ahtng/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-25T11:35:09.304408341+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmp3vlii6qe/port-20653.txt","pid":20653,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T11:35:09.30601+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":20653} +{"time":"2025-12-25T11:35:09.30594987+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-20653-20769-3087691760/socket","Net":"unix"}} +{"time":"2025-12-25T11:35:09.369467149+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T11:35:09.375606365+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"605ahtng","id":"1(@)"} +{"time":"2025-12-25T11:35:09.926496776+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"605ahtng","id":"1(@)"} +{"time":"2025-12-25T11:35:10.167659949+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T11:35:10.167754484+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T11:35:10.167807689+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T11:35:10.167821801+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T11:35:10.167919924+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-20653-20769-3087691760/socket","Net":"unix"}} +{"time":"2025-12-25T11:35:10.224907747+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-25T11:35:10.224935225+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-25T11:35:10.224954175+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/run-20251225_113509-605ahtng/logs/debug-internal.log b/wandb/run-20251225_113509-605ahtng/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..a2a25067b6b7a047d60a88a2b3d3e943496bdb95 --- /dev/null +++ b/wandb/run-20251225_113509-605ahtng/logs/debug-internal.log @@ -0,0 +1,15 @@ +{"time":"2025-12-25T11:35:09.375701438+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T11:35:09.926285637+07:00","level":"INFO","msg":"stream: created new stream","id":"605ahtng"} +{"time":"2025-12-25T11:35:09.926489479+07:00","level":"INFO","msg":"stream: started","id":"605ahtng"} +{"time":"2025-12-25T11:35:09.929864536+07:00","level":"INFO","msg":"sender: started","stream_id":"605ahtng"} +{"time":"2025-12-25T11:35:09.929942154+07:00","level":"INFO","msg":"handler: started","stream_id":"605ahtng"} +{"time":"2025-12-25T11:35:09.929915373+07:00","level":"INFO","msg":"writer: started","stream_id":"605ahtng"} +{"time":"2025-12-25T11:35:10.156405445+07:00","level":"ERROR","msg":"HTTP error","status":401,"method":"POST","url":"https://api.wandb.ai/graphql"} +{"time":"2025-12-25T11:35:10.156545252+07:00","level":"ERROR","msg":"runupserter: failed to init run","error":"returned error 401: {\"data\":{\"upsertBucket\":null},\"errors\":[{\"message\":\"user is not logged in\",\"path\":[\"upsertBucket\"],\"extensions\":{\"code\":\"PERMISSION_ERROR\"}}]}"} +{"time":"2025-12-25T11:35:10.167733434+07:00","level":"INFO","msg":"stream: closing","id":"605ahtng"} +{"time":"2025-12-25T11:35:10.168204712+07:00","level":"ERROR","msg":"sender: uploadConfigFile: runupserter: run not yet initialized"} +{"time":"2025-12-25T11:35:10.218814625+07:00","level":"ERROR","msg":"runfiles: runupserter: run not yet initialized"} +{"time":"2025-12-25T11:35:10.221619936+07:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"} +{"time":"2025-12-25T11:35:10.221674069+07:00","level":"INFO","msg":"handler: closed","stream_id":"605ahtng"} +{"time":"2025-12-25T11:35:10.221727328+07:00","level":"INFO","msg":"sender: closed","stream_id":"605ahtng"} +{"time":"2025-12-25T11:35:10.221733424+07:00","level":"INFO","msg":"stream: closed","id":"605ahtng"} diff --git a/wandb/run-20251225_113509-605ahtng/logs/debug.log b/wandb/run-20251225_113509-605ahtng/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..1bb49a6f411e564d5b18f2379ef75b23c9a893e8 --- /dev/null +++ b/wandb/run-20251225_113509-605ahtng/logs/debug.log @@ -0,0 +1,15 @@ +2025-12-25 11:35:09,024 INFO MainThread:20653 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 11:35:09,024 INFO MainThread:20653 [wandb_setup.py:_flush():80] Configure stats pid to 20653 +2025-12-25 11:35:09,024 INFO MainThread:20653 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 11:35:09,024 INFO MainThread:20653 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 11:35:09,024 INFO MainThread:20653 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 11:35:09,024 INFO MainThread:20653 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_113509-605ahtng/logs/debug.log +2025-12-25 11:35:09,024 INFO MainThread:20653 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_113509-605ahtng/logs/debug-internal.log +2025-12-25 11:35:09,024 INFO MainThread:20653 [wandb_init.py:init():840] calling init triggers +2025-12-25 11:35:09,024 INFO MainThread:20653 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 11:35:09,024 INFO MainThread:20653 [wandb_init.py:init():888] starting backend +2025-12-25 11:35:09,368 INFO MainThread:20653 [wandb_init.py:init():891] sending inform_init request +2025-12-25 11:35:09,373 INFO MainThread:20653 [wandb_init.py:init():899] backend started and connected +2025-12-25 11:35:09,373 INFO MainThread:20653 [wandb_init.py:init():969] updated telemetry +2025-12-25 11:35:09,376 INFO MainThread:20653 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout diff --git a/wandb/run-20251225_113509-605ahtng/run-605ahtng.wandb b/wandb/run-20251225_113509-605ahtng/run-605ahtng.wandb new file mode 100644 index 0000000000000000000000000000000000000000..e9f6593bf6fc081efad853e7954be2d0b8375e84 Binary files /dev/null and b/wandb/run-20251225_113509-605ahtng/run-605ahtng.wandb differ diff --git a/wandb/run-20251225_114329-5vxk2v1x/files/code/src/imagenet/train_model.py b/wandb/run-20251225_114329-5vxk2v1x/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..ad886085b66d9b18d44a70fa31214a8ca254d6ee --- /dev/null +++ b/wandb/run-20251225_114329-5vxk2v1x/files/code/src/imagenet/train_model.py @@ -0,0 +1,262 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251225_114329-5vxk2v1x/files/config.yaml b/wandb/run-20251225_114329-5vxk2v1x/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..45bac8356469e0dd73f8192c8b2a28e5af8d5edd --- /dev/null +++ b/wandb/run-20251225_114329-5vxk2v1x/files/config.yaml @@ -0,0 +1,111 @@ +_wandb: + value: + cli_version: 0.23.0 + code_path: code/src/imagenet/train_model.py + e: + 30kwis5h5n4dkeytog60vt9881gt6hm8: + args: + - --input-size + - "32" + - --data-set + - CIFAR + - --patch-size + - "4" + - --hidden-size + - "128" + - --num-hidden-layers + - "6" + - --num-attention-heads + - "4" + - --intermediate-size + - "512" + - --position-embeddings + - learnable + - --lr + - "5e-3" + - --epochs + - "100" + - --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 + - /mnt/data/vinhbk/weights/lmc/cifar100 + - --data-path + - /mnt/data/vinhbk/datasets/cifar100 + codePath: src/imagenet/train_model.py + codePathLocal: src/imagenet/train_model.py + cpu_count: 24 + cpu_count_logical: 24 + cudaVersion: "13.0" + disk: + /: + total: "1081101176832" + used: "67539881984" + email: khanhvinhbui0512@gmail.com + executable: /home/khanhvinh9/miniconda3/envs/lmc/bin/python + git: + commit: 503490314d9f0bef18a30ec7957b3949bd36c751 + remote: https://github.com/MLResearchX/lmc-transformer.git + gpu: NVIDIA GeForce RTX 5060 Laptop GPU + gpu_count: 1 + gpu_nvidia: + - architecture: Blackwell + memoryTotal: "8546942976" + name: NVIDIA GeForce RTX 5060 Laptop GPU + uuid: GPU-6363d584-7737-137f-cc48-e39ef1a9efbd + host: DESKTOP-UCJ5IKF + memory: + total: "16466997248" + os: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 + program: /home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py + python: CPython 3.10.19 + root: /home/khanhvinh9/lmc-transformer + startedAt: "2025-12-25T04:43:29.013900Z" + writerId: 30kwis5h5n4dkeytog60vt9881gt6hm8 + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "2": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "3": + - 13 + "4": 3.10.19 + "5": 0.23.0 + "6": 4.56.1 + "8": + - 2 + "12": 0.23.0 + "13": linux-x86_64 diff --git a/wandb/run-20251225_114329-5vxk2v1x/files/diff.patch b/wandb/run-20251225_114329-5vxk2v1x/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..64066b87ba66f006923a1de989a5c15561cbb8ff --- /dev/null +++ b/wandb/run-20251225_114329-5vxk2v1x/files/diff.patch @@ -0,0 +1,134 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..b4f18b4 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..f57f9d2 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': + dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..ad88608 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,13 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +212,7 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_114329-5vxk2v1x/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251225_114329-5vxk2v1x/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..64066b87ba66f006923a1de989a5c15561cbb8ff --- /dev/null +++ b/wandb/run-20251225_114329-5vxk2v1x/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,134 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..b4f18b4 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..f57f9d2 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': + dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..ad88608 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,13 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +212,7 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_114329-5vxk2v1x/files/output.log b/wandb/run-20251225_114329-5vxk2v1x/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..0b1f15fb4a2161ac442e020d349c3add8f2b1c3b --- /dev/null +++ b/wandb/run-20251225_114329-5vxk2v1x/files/output.log @@ -0,0 +1,24 @@ +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 262, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 87, in main + train_loader, val_loader = imagenet_data_loader(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 42, in imagenet_data_loader + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/datasets.py", line 63, in build_dataset + dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torchvision/datasets/cifar.py", line 69, in __init__ + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") +RuntimeError: Dataset not found or corrupted. You can use download=True to download it +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 262, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 87, in main + train_loader, val_loader = imagenet_data_loader(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 42, in imagenet_data_loader + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/datasets.py", line 63, in build_dataset + dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torchvision/datasets/cifar.py", line 69, in __init__ + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") +RuntimeError: Dataset not found or corrupted. You can use download=True to download it diff --git a/wandb/run-20251225_114329-5vxk2v1x/files/requirements.txt b/wandb/run-20251225_114329-5vxk2v1x/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251225_114329-5vxk2v1x/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251225_114329-5vxk2v1x/files/wandb-metadata.json b/wandb/run-20251225_114329-5vxk2v1x/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..594c13f0062ffa0c0708cec95949ac7641f5405c --- /dev/null +++ b/wandb/run-20251225_114329-5vxk2v1x/files/wandb-metadata.json @@ -0,0 +1,81 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-25T04:43:29.013900Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "128", + "--num-hidden-layers", + "6", + "--num-attention-heads", + "4", + "--intermediate-size", + "512", + "--position-embeddings", + "learnable", + "--lr", + "5e-3", + "--epochs", + "100", + "--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", + "/mnt/data/vinhbk/weights/lmc/cifar100", + "--data-path", + "/mnt/data/vinhbk/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67539881984" + } + }, + "memory": { + "total": "16466997248" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "30kwis5h5n4dkeytog60vt9881gt6hm8" +} \ No newline at end of file diff --git a/wandb/run-20251225_114329-5vxk2v1x/files/wandb-summary.json b/wandb/run-20251225_114329-5vxk2v1x/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..794feb4b9cafbfea9c76abd3f919b15557e9d45f --- /dev/null +++ b/wandb/run-20251225_114329-5vxk2v1x/files/wandb-summary.json @@ -0,0 +1 @@ +{"_runtime":1,"_wandb":{"runtime":1}} \ No newline at end of file diff --git a/wandb/run-20251225_114329-5vxk2v1x/logs/debug-core.log b/wandb/run-20251225_114329-5vxk2v1x/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..cbc80165ea13bc430b64b08d0372bebba44c7010 --- /dev/null +++ b/wandb/run-20251225_114329-5vxk2v1x/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-25T11:43:29.295295505+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmpfkk2g06e/port-23995.txt","pid":23995,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T11:43:29.297544629+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":23995} +{"time":"2025-12-25T11:43:29.297552436+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-23995-24138-2203404552/socket","Net":"unix"}} +{"time":"2025-12-25T11:43:29.369356009+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T11:43:29.381886375+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"5vxk2v1x","id":"1(@)"} +{"time":"2025-12-25T11:43:30.011610384+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"5vxk2v1x","id":"1(@)"} +{"time":"2025-12-25T11:43:31.797450317+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T11:43:31.797521041+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T11:43:31.797563658+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T11:43:31.797562308+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T11:43:31.797622978+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-23995-24138-2203404552/socket","Net":"unix"}} +{"time":"2025-12-25T11:43:33.555843413+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-25T11:43:33.55586973+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-25T11:43:33.555885133+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/run-20251225_114329-5vxk2v1x/logs/debug-internal.log b/wandb/run-20251225_114329-5vxk2v1x/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..7eef428ac5659c0852e31e4f854e3ee3ca8ce5d1 --- /dev/null +++ b/wandb/run-20251225_114329-5vxk2v1x/logs/debug-internal.log @@ -0,0 +1,11 @@ +{"time":"2025-12-25T11:43:29.382028582+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T11:43:30.011346796+07:00","level":"INFO","msg":"stream: created new stream","id":"5vxk2v1x"} +{"time":"2025-12-25T11:43:30.011482695+07:00","level":"INFO","msg":"handler: started","stream_id":"5vxk2v1x"} +{"time":"2025-12-25T11:43:30.011601479+07:00","level":"INFO","msg":"stream: started","id":"5vxk2v1x"} +{"time":"2025-12-25T11:43:30.011631717+07:00","level":"INFO","msg":"writer: started","stream_id":"5vxk2v1x"} +{"time":"2025-12-25T11:43:30.011666923+07:00","level":"INFO","msg":"sender: started","stream_id":"5vxk2v1x"} +{"time":"2025-12-25T11:43:31.797514947+07:00","level":"INFO","msg":"stream: closing","id":"5vxk2v1x"} +{"time":"2025-12-25T11:43:33.150053118+07:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"} +{"time":"2025-12-25T11:43:33.55012408+07:00","level":"INFO","msg":"handler: closed","stream_id":"5vxk2v1x"} +{"time":"2025-12-25T11:43:33.551129059+07:00","level":"INFO","msg":"sender: closed","stream_id":"5vxk2v1x"} +{"time":"2025-12-25T11:43:33.551146327+07:00","level":"INFO","msg":"stream: closed","id":"5vxk2v1x"} diff --git a/wandb/run-20251225_114329-5vxk2v1x/logs/debug.log b/wandb/run-20251225_114329-5vxk2v1x/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..8b7c7f498b5195a2a0c263cbd6daf8af9a0d0d68 --- /dev/null +++ b/wandb/run-20251225_114329-5vxk2v1x/logs/debug.log @@ -0,0 +1,23 @@ +2025-12-25 11:43:29,018 INFO MainThread:23995 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 11:43:29,018 INFO MainThread:23995 [wandb_setup.py:_flush():80] Configure stats pid to 23995 +2025-12-25 11:43:29,018 INFO MainThread:23995 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 11:43:29,018 INFO MainThread:23995 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 11:43:29,018 INFO MainThread:23995 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 11:43:29,018 INFO MainThread:23995 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_114329-5vxk2v1x/logs/debug.log +2025-12-25 11:43:29,018 INFO MainThread:23995 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_114329-5vxk2v1x/logs/debug-internal.log +2025-12-25 11:43:29,018 INFO MainThread:23995 [wandb_init.py:init():840] calling init triggers +2025-12-25 11:43:29,018 INFO MainThread:23995 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 11:43:29,018 INFO MainThread:23995 [wandb_init.py:init():888] starting backend +2025-12-25 11:43:29,368 INFO MainThread:23995 [wandb_init.py:init():891] sending inform_init request +2025-12-25 11:43:29,375 INFO MainThread:23995 [wandb_init.py:init():899] backend started and connected +2025-12-25 11:43:29,377 INFO MainThread:23995 [wandb_init.py:init():969] updated telemetry +2025-12-25 11:43:29,381 INFO MainThread:23995 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 11:43:30,398 INFO MainThread:23995 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 11:43:30,631 INFO MainThread:23995 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 11:43:30,632 INFO MainThread:23995 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 11:43:30,633 INFO MainThread:23995 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 11:43:30,633 INFO MainThread:23995 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 11:43:30,636 INFO MainThread:23995 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-25 11:43:31,797 INFO wandb-AsyncioManager-main:23995 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-25 11:43:31,797 INFO wandb-AsyncioManager-main:23995 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. diff --git a/wandb/run-20251225_114329-5vxk2v1x/run-5vxk2v1x.wandb b/wandb/run-20251225_114329-5vxk2v1x/run-5vxk2v1x.wandb new file mode 100644 index 0000000000000000000000000000000000000000..349bae92267e8ee884af60b9d20c36ba8835cfb2 Binary files /dev/null and b/wandb/run-20251225_114329-5vxk2v1x/run-5vxk2v1x.wandb differ diff --git a/wandb/run-20251225_115532-nv5va7ak/files/code/src/imagenet/train_model.py b/wandb/run-20251225_115532-nv5va7ak/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..ad886085b66d9b18d44a70fa31214a8ca254d6ee --- /dev/null +++ b/wandb/run-20251225_115532-nv5va7ak/files/code/src/imagenet/train_model.py @@ -0,0 +1,262 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251225_115532-nv5va7ak/files/config.yaml b/wandb/run-20251225_115532-nv5va7ak/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..804498a144c0d50cb13b3bc1aba943f58eacbd61 --- /dev/null +++ b/wandb/run-20251225_115532-nv5va7ak/files/config.yaml @@ -0,0 +1,111 @@ +_wandb: + value: + cli_version: 0.23.0 + code_path: code/src/imagenet/train_model.py + e: + eyx2bkvmmc5xmglleqhommvc2b1qz6bd: + args: + - --input-size + - "32" + - --data-set + - CIFAR + - --patch-size + - "4" + - --hidden-size + - "128" + - --num-hidden-layers + - "6" + - --num-attention-heads + - "4" + - --intermediate-size + - "512" + - --position-embeddings + - learnable + - --lr + - "5e-3" + - --epochs + - "100" + - --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 + - /mnt/d/NCKH/weights/lmc/cifar100 + - --data-path + - /mnt/d/NCKH/datasets/cifar100 + codePath: src/imagenet/train_model.py + codePathLocal: src/imagenet/train_model.py + cpu_count: 24 + cpu_count_logical: 24 + cudaVersion: "13.0" + disk: + /: + total: "1081101176832" + used: "67540508672" + email: khanhvinhbui0512@gmail.com + executable: /home/khanhvinh9/miniconda3/envs/lmc/bin/python + git: + commit: 503490314d9f0bef18a30ec7957b3949bd36c751 + remote: https://github.com/MLResearchX/lmc-transformer.git + gpu: NVIDIA GeForce RTX 5060 Laptop GPU + gpu_count: 1 + gpu_nvidia: + - architecture: Blackwell + memoryTotal: "8546942976" + name: NVIDIA GeForce RTX 5060 Laptop GPU + uuid: GPU-6363d584-7737-137f-cc48-e39ef1a9efbd + host: DESKTOP-UCJ5IKF + memory: + total: "16466997248" + os: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 + program: /home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py + python: CPython 3.10.19 + root: /home/khanhvinh9/lmc-transformer + startedAt: "2025-12-25T04:55:32.335184Z" + writerId: eyx2bkvmmc5xmglleqhommvc2b1qz6bd + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "2": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "3": + - 13 + "4": 3.10.19 + "5": 0.23.0 + "6": 4.56.1 + "8": + - 2 + "12": 0.23.0 + "13": linux-x86_64 diff --git a/wandb/run-20251225_115532-nv5va7ak/files/diff.patch b/wandb/run-20251225_115532-nv5va7ak/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..0949c1f0dce1fb3d28b8050d5983a6a24963cd62 --- /dev/null +++ b/wandb/run-20251225_115532-nv5va7ak/files/diff.patch @@ -0,0 +1,135 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..ad88608 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,13 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +212,7 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_115532-nv5va7ak/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251225_115532-nv5va7ak/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..0949c1f0dce1fb3d28b8050d5983a6a24963cd62 --- /dev/null +++ b/wandb/run-20251225_115532-nv5va7ak/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,135 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..ad88608 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,13 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +212,7 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_115532-nv5va7ak/files/output.log b/wandb/run-20251225_115532-nv5va7ak/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..261fbaabbe12c2c77e4ad36cf3ee43dec8558ba4 --- /dev/null +++ b/wandb/run-20251225_115532-nv5va7ak/files/output.log @@ -0,0 +1,53 @@ +TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We recommend migrating to PyTorch classes or pinning your version of Transformers. +Starting training... +JAX devices: [CudaDevice(id=0)] +Using 1 devices +Epoch 0 +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 262, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 165, in main + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 494, in __iter__ + return self._get_iterator() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 427, in _get_iterator + return _MultiProcessingDataLoaderIter(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1170, in __init__ + w.start() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/process.py", line 121, in start + self._popen = self._Popen(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 224, in _Popen + return _default_context.get_context().Process._Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 288, in _Popen + return Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__ + super().__init__(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__ + self._launch(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 62, in _launch + f.write(fp.getbuffer()) +KeyboardInterrupt +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 262, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 165, in main + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 494, in __iter__ + return self._get_iterator() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 427, in _get_iterator + return _MultiProcessingDataLoaderIter(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1170, in __init__ + w.start() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/process.py", line 121, in start + self._popen = self._Popen(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 224, in _Popen + return _default_context.get_context().Process._Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 288, in _Popen + return Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__ + super().__init__(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__ + self._launch(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 62, in _launch + f.write(fp.getbuffer()) +KeyboardInterrupt diff --git a/wandb/run-20251225_115532-nv5va7ak/files/requirements.txt b/wandb/run-20251225_115532-nv5va7ak/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251225_115532-nv5va7ak/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251225_115532-nv5va7ak/files/wandb-metadata.json b/wandb/run-20251225_115532-nv5va7ak/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..94b2781d04de62ea0e203889281d5cea11c72606 --- /dev/null +++ b/wandb/run-20251225_115532-nv5va7ak/files/wandb-metadata.json @@ -0,0 +1,81 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-25T04:55:32.335184Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "128", + "--num-hidden-layers", + "6", + "--num-attention-heads", + "4", + "--intermediate-size", + "512", + "--position-embeddings", + "learnable", + "--lr", + "5e-3", + "--epochs", + "100", + "--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", + "/mnt/d/NCKH/weights/lmc/cifar100", + "--data-path", + "/mnt/d/NCKH/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67540508672" + } + }, + "memory": { + "total": "16466997248" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "eyx2bkvmmc5xmglleqhommvc2b1qz6bd" +} \ No newline at end of file diff --git a/wandb/run-20251225_115532-nv5va7ak/files/wandb-summary.json b/wandb/run-20251225_115532-nv5va7ak/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..8cd5e3ef973e8a094ab6119f85e5f0e91f159782 --- /dev/null +++ b/wandb/run-20251225_115532-nv5va7ak/files/wandb-summary.json @@ -0,0 +1 @@ +{"_wandb":{"runtime":33},"_runtime":33} \ No newline at end of file diff --git a/wandb/run-20251225_115532-nv5va7ak/logs/debug-core.log b/wandb/run-20251225_115532-nv5va7ak/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..4c2631bc13c28c1a87f73d49dd96847f0833a5e0 --- /dev/null +++ b/wandb/run-20251225_115532-nv5va7ak/logs/debug-core.log @@ -0,0 +1,12 @@ +{"time":"2025-12-25T11:55:32.441966402+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmpqqilnf73/port-36170.txt","pid":36170,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T11:55:32.4430092+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":36170} +{"time":"2025-12-25T11:55:32.442963357+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-36170-36280-3895705256/socket","Net":"unix"}} +{"time":"2025-12-25T11:55:32.6275177+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T11:55:32.635364546+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"nv5va7ak","id":"1(@)"} +{"time":"2025-12-25T11:55:33.119886595+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"nv5va7ak","id":"1(@)"} +{"time":"2025-12-25T11:56:06.025913212+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T11:56:06.026148716+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T11:56:06.02612658+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T11:56:06.026245818+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-36170-36280-3895705256/socket","Net":"unix"}} +{"time":"2025-12-25T11:56:06.026298636+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T11:56:06.851567683+07:00","level":"INFO","msg":"server: parent process exited, terminating service process"} diff --git a/wandb/run-20251225_115532-nv5va7ak/logs/debug-internal.log b/wandb/run-20251225_115532-nv5va7ak/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..d8704569fff46907003565211a8da1ea1f1d8c2c --- /dev/null +++ b/wandb/run-20251225_115532-nv5va7ak/logs/debug-internal.log @@ -0,0 +1,8 @@ +{"time":"2025-12-25T11:55:32.635506289+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T11:55:33.119566136+07:00","level":"INFO","msg":"stream: created new stream","id":"nv5va7ak"} +{"time":"2025-12-25T11:55:33.119692659+07:00","level":"INFO","msg":"handler: started","stream_id":"nv5va7ak"} +{"time":"2025-12-25T11:55:33.119876248+07:00","level":"INFO","msg":"stream: started","id":"nv5va7ak"} +{"time":"2025-12-25T11:55:33.119909315+07:00","level":"INFO","msg":"writer: started","stream_id":"nv5va7ak"} +{"time":"2025-12-25T11:55:33.119933002+07:00","level":"INFO","msg":"sender: started","stream_id":"nv5va7ak"} +{"time":"2025-12-25T11:56:06.026061654+07:00","level":"INFO","msg":"stream: closing","id":"nv5va7ak"} +{"time":"2025-12-25T11:56:06.723283342+07:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"} diff --git a/wandb/run-20251225_115532-nv5va7ak/logs/debug.log b/wandb/run-20251225_115532-nv5va7ak/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..2b252c830fdcb5272bb9e194dfdcb045efb025a4 --- /dev/null +++ b/wandb/run-20251225_115532-nv5va7ak/logs/debug.log @@ -0,0 +1,790 @@ +2025-12-25 11:55:32,337 INFO MainThread:36170 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 11:55:32,337 INFO MainThread:36170 [wandb_setup.py:_flush():80] Configure stats pid to 36170 +2025-12-25 11:55:32,338 INFO MainThread:36170 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 11:55:32,338 INFO MainThread:36170 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 11:55:32,338 INFO MainThread:36170 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 11:55:32,338 INFO MainThread:36170 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_115532-nv5va7ak/logs/debug.log +2025-12-25 11:55:32,338 INFO MainThread:36170 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_115532-nv5va7ak/logs/debug-internal.log +2025-12-25 11:55:32,338 INFO MainThread:36170 [wandb_init.py:init():840] calling init triggers +2025-12-25 11:55:32,338 INFO MainThread:36170 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 11:55:32,338 INFO MainThread:36170 [wandb_init.py:init():888] starting backend +2025-12-25 11:55:32,627 INFO MainThread:36170 [wandb_init.py:init():891] sending inform_init request +2025-12-25 11:55:32,632 INFO MainThread:36170 [wandb_init.py:init():899] backend started and connected +2025-12-25 11:55:32,634 INFO MainThread:36170 [wandb_init.py:init():969] updated telemetry +2025-12-25 11:55:32,639 INFO MainThread:36170 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 11:55:33,620 INFO MainThread:36170 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 11:55:33,733 INFO MainThread:36170 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 11:55:33,734 INFO MainThread:36170 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 11:55:33,734 INFO MainThread:36170 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 11:55:33,734 INFO MainThread:36170 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 11:55:33,736 INFO MainThread:36170 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-25 11:56:06,025 INFO wandb-AsyncioManager-main:36170 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-25 11:56:06,026 INFO wandb-AsyncioManager-main:36170 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. +2025-12-25 11:56:06,439 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,442 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,443 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,443 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,443 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,456 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,457 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,457 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,457 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,458 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,458 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,459 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,459 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,459 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,459 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,459 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,459 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,460 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,460 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,460 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,460 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,460 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,461 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,461 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,546 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,559 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,559 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,559 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,560 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,565 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,566 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,566 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,566 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,566 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,566 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,567 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,567 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,567 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,567 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,567 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,567 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,568 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,568 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,568 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,568 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,568 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,568 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,569 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,569 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,569 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,569 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,569 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,569 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,569 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,570 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,570 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,570 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,618 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 11:56:06,618 ERROR wandb-AsyncioManager-main:36170 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost diff --git a/wandb/run-20251225_115532-nv5va7ak/run-nv5va7ak.wandb b/wandb/run-20251225_115532-nv5va7ak/run-nv5va7ak.wandb new file mode 100644 index 0000000000000000000000000000000000000000..30977402b08eb894463b83d7de7a007255879a32 Binary files /dev/null and b/wandb/run-20251225_115532-nv5va7ak/run-nv5va7ak.wandb differ diff --git a/wandb/run-20251225_120039-17r3f9k1/files/code/src/imagenet/train_model.py b/wandb/run-20251225_120039-17r3f9k1/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c48bfefd7467518c68801e1068bea786a32c1e92 --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/files/code/src/imagenet/train_model.py @@ -0,0 +1,264 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251225_120039-17r3f9k1/files/config.yaml b/wandb/run-20251225_120039-17r3f9k1/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..487e61a88bed9544e10bc4222e94f465c7dbccc6 --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/files/config.yaml @@ -0,0 +1,112 @@ +_wandb: + value: + cli_version: 0.23.0 + code_path: code/src/imagenet/train_model.py + e: + j3ccj5pt0vxvh9he9wxiavmchk2a4a67: + args: + - --input-size + - "32" + - --data-set + - CIFAR + - --patch-size + - "4" + - --hidden-size + - "128" + - --num-hidden-layers + - "6" + - --num-attention-heads + - "4" + - --intermediate-size + - "512" + - --position-embeddings + - learnable + - --lr + - "5e-3" + - --epochs + - "100" + - --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 + - /mnt/d/NCKH/weights/lmc/cifar100 + - --data-path + - /mnt/d/NCKH/datasets/cifar100 + codePath: src/imagenet/train_model.py + codePathLocal: src/imagenet/train_model.py + cpu_count: 24 + cpu_count_logical: 24 + cudaVersion: "13.0" + disk: + /: + total: "1081101176832" + used: "67540955136" + email: khanhvinhbui0512@gmail.com + executable: /home/khanhvinh9/miniconda3/envs/lmc/bin/python + git: + commit: 503490314d9f0bef18a30ec7957b3949bd36c751 + remote: https://github.com/MLResearchX/lmc-transformer.git + gpu: NVIDIA GeForce RTX 5060 Laptop GPU + gpu_count: 1 + gpu_nvidia: + - architecture: Blackwell + memoryTotal: "8546942976" + name: NVIDIA GeForce RTX 5060 Laptop GPU + uuid: GPU-6363d584-7737-137f-cc48-e39ef1a9efbd + host: DESKTOP-UCJ5IKF + memory: + total: "16466997248" + os: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 + program: /home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py + python: CPython 3.10.19 + root: /home/khanhvinh9/lmc-transformer + startedAt: "2025-12-25T05:00:39.962187Z" + writerId: j3ccj5pt0vxvh9he9wxiavmchk2a4a67 + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "2": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "3": + - 13 + - 61 + "4": 3.10.19 + "5": 0.23.0 + "6": 4.56.1 + "8": + - 2 + "12": 0.23.0 + "13": linux-x86_64 diff --git a/wandb/run-20251225_120039-17r3f9k1/files/diff.patch b/wandb/run-20251225_120039-17r3f9k1/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..00762339a061f1e7187065f66f071e30db8800f7 --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/files/diff.patch @@ -0,0 +1,137 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_120039-17r3f9k1/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251225_120039-17r3f9k1/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..00762339a061f1e7187065f66f071e30db8800f7 --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,137 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_120039-17r3f9k1/files/output.log b/wandb/run-20251225_120039-17r3f9k1/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..40b449aa0cf2b69250e1fd45e78dcb5251ac1cfc --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/files/output.log @@ -0,0 +1,27 @@ +TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We recommend migrating to PyTorch classes or pinning your version of Transformers. +Starting training... +JAX devices: [CudaDevice(id=0)] +Using 1 devices +Epoch 0 +WARNING:absl:[process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 1 +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 207, in main + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/flax/training/checkpoints.py", line 693, in save_checkpoint + orbax_checkpointer.save( + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 254, in save + raise ValueError(f'Destination {directory} already exists.') +ValueError: Destination /mnt/d/NCKH/weights/lmc/cifar100/lr0.005-learnable-epochs100-batch128-seed0/last_780 already exists. +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 207, in main + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/flax/training/checkpoints.py", line 693, in save_checkpoint + orbax_checkpointer.save( + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 254, in save + raise ValueError(f'Destination {directory} already exists.') +ValueError: Destination /mnt/d/NCKH/weights/lmc/cifar100/lr0.005-learnable-epochs100-batch128-seed0/last_780 already exists. diff --git a/wandb/run-20251225_120039-17r3f9k1/files/requirements.txt b/wandb/run-20251225_120039-17r3f9k1/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251225_120039-17r3f9k1/files/wandb-metadata.json b/wandb/run-20251225_120039-17r3f9k1/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..b437d6015cf8459a51aa208b47b52f80f3f40d5b --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/files/wandb-metadata.json @@ -0,0 +1,81 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-25T05:00:39.962187Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "128", + "--num-hidden-layers", + "6", + "--num-attention-heads", + "4", + "--intermediate-size", + "512", + "--position-embeddings", + "learnable", + "--lr", + "5e-3", + "--epochs", + "100", + "--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", + "/mnt/d/NCKH/weights/lmc/cifar100", + "--data-path", + "/mnt/d/NCKH/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67540955136" + } + }, + "memory": { + "total": "16466997248" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "j3ccj5pt0vxvh9he9wxiavmchk2a4a67" +} \ No newline at end of file diff --git a/wandb/run-20251225_120039-17r3f9k1/files/wandb-summary.json b/wandb/run-20251225_120039-17r3f9k1/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..11370c4ea9ed9e0600f7c2eb8027fa2b616b9c59 --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/files/wandb-summary.json @@ -0,0 +1 @@ +{"_timestamp":1.7666390698490946e+09,"val/acc5":9.375,"val/loss":4.59375,"_runtime":239.236678129,"train/loss":4.5625,"train/acc5":12.5,"lr":0.0017955126240849495,"_step":780,"_wandb":{"runtime":239},"val/acc1":2.34375,"train/acc1":2.34375} \ No newline at end of file diff --git a/wandb/run-20251225_120039-17r3f9k1/logs/debug-core.log b/wandb/run-20251225_120039-17r3f9k1/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..1dc4ecbccb4c7804ed50c2c43d8776e777c52269 --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-25T12:00:40.065959107+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmpe_gy0fqb/port-38983.txt","pid":38983,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T12:00:40.067091997+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":38983} +{"time":"2025-12-25T12:00:40.067010125+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-38983-39092-3896389027/socket","Net":"unix"}} +{"time":"2025-12-25T12:00:40.243849484+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T12:00:40.251764569+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"17r3f9k1","id":"1(@)"} +{"time":"2025-12-25T12:00:40.713480929+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"17r3f9k1","id":"1(@)"} +{"time":"2025-12-25T12:04:29.887351392+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T12:04:29.887426506+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T12:04:29.887478226+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T12:04:29.887446371+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T12:04:29.887550367+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-38983-39092-3896389027/socket","Net":"unix"}} +{"time":"2025-12-25T12:04:30.97412261+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-25T12:04:30.974154747+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-25T12:04:30.974164149+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/run-20251225_120039-17r3f9k1/logs/debug-internal.log b/wandb/run-20251225_120039-17r3f9k1/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..7f9582735030816693be5344f2453f013ab8f9fe --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/logs/debug-internal.log @@ -0,0 +1,11 @@ +{"time":"2025-12-25T12:00:40.251901299+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T12:00:40.713208051+07:00","level":"INFO","msg":"stream: created new stream","id":"17r3f9k1"} +{"time":"2025-12-25T12:00:40.713382585+07:00","level":"INFO","msg":"handler: started","stream_id":"17r3f9k1"} +{"time":"2025-12-25T12:00:40.713471773+07:00","level":"INFO","msg":"stream: started","id":"17r3f9k1"} +{"time":"2025-12-25T12:00:40.713509747+07:00","level":"INFO","msg":"writer: started","stream_id":"17r3f9k1"} +{"time":"2025-12-25T12:00:40.71357019+07:00","level":"INFO","msg":"sender: started","stream_id":"17r3f9k1"} +{"time":"2025-12-25T12:04:29.887392463+07:00","level":"INFO","msg":"stream: closing","id":"17r3f9k1"} +{"time":"2025-12-25T12:04:30.553184702+07:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"} +{"time":"2025-12-25T12:04:30.968568731+07:00","level":"INFO","msg":"handler: closed","stream_id":"17r3f9k1"} +{"time":"2025-12-25T12:04:30.968765344+07:00","level":"INFO","msg":"sender: closed","stream_id":"17r3f9k1"} +{"time":"2025-12-25T12:04:30.968777976+07:00","level":"INFO","msg":"stream: closed","id":"17r3f9k1"} diff --git a/wandb/run-20251225_120039-17r3f9k1/logs/debug.log b/wandb/run-20251225_120039-17r3f9k1/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..69632f6f9fa124e04d1aaf1740c15b28c1ef42d3 --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/logs/debug.log @@ -0,0 +1,23 @@ +2025-12-25 12:00:39,967 INFO MainThread:38983 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 12:00:39,967 INFO MainThread:38983 [wandb_setup.py:_flush():80] Configure stats pid to 38983 +2025-12-25 12:00:39,967 INFO MainThread:38983 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 12:00:39,967 INFO MainThread:38983 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 12:00:39,967 INFO MainThread:38983 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 12:00:39,967 INFO MainThread:38983 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_120039-17r3f9k1/logs/debug.log +2025-12-25 12:00:39,967 INFO MainThread:38983 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_120039-17r3f9k1/logs/debug-internal.log +2025-12-25 12:00:39,967 INFO MainThread:38983 [wandb_init.py:init():840] calling init triggers +2025-12-25 12:00:39,967 INFO MainThread:38983 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 12:00:39,967 INFO MainThread:38983 [wandb_init.py:init():888] starting backend +2025-12-25 12:00:40,244 INFO MainThread:38983 [wandb_init.py:init():891] sending inform_init request +2025-12-25 12:00:40,249 INFO MainThread:38983 [wandb_init.py:init():899] backend started and connected +2025-12-25 12:00:40,250 INFO MainThread:38983 [wandb_init.py:init():969] updated telemetry +2025-12-25 12:00:40,255 INFO MainThread:38983 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 12:00:41,136 INFO MainThread:38983 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 12:00:41,252 INFO MainThread:38983 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 12:00:41,252 INFO MainThread:38983 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 12:00:41,252 INFO MainThread:38983 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 12:00:41,252 INFO MainThread:38983 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 12:00:41,254 INFO MainThread:38983 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-25 12:04:29,887 INFO wandb-AsyncioManager-main:38983 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-25 12:04:29,887 INFO wandb-AsyncioManager-main:38983 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. diff --git a/wandb/run-20251225_120039-17r3f9k1/run-17r3f9k1.wandb b/wandb/run-20251225_120039-17r3f9k1/run-17r3f9k1.wandb new file mode 100644 index 0000000000000000000000000000000000000000..0e477ec161df581fe047a9f1461d48f6ff28adf9 --- /dev/null +++ b/wandb/run-20251225_120039-17r3f9k1/run-17r3f9k1.wandb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b6168a6b143df1f2ac192d6c7287fcf53cedcd598886d5d2a00f9a61295e6761 +size 163227 diff --git a/wandb/run-20251225_121621-23cpikjp/files/code/src/imagenet/train_model.py b/wandb/run-20251225_121621-23cpikjp/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c48bfefd7467518c68801e1068bea786a32c1e92 --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/files/code/src/imagenet/train_model.py @@ -0,0 +1,264 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251225_121621-23cpikjp/files/config.yaml b/wandb/run-20251225_121621-23cpikjp/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f06a79077ac550b746c01d14f6b9d277178e1caf --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/files/config.yaml @@ -0,0 +1,112 @@ +_wandb: + value: + cli_version: 0.23.0 + code_path: code/src/imagenet/train_model.py + e: + gxr0tuqothrxfogjmk2mbhag5kge133z: + args: + - --input-size + - "32" + - --data-set + - CIFAR + - --patch-size + - "4" + - --hidden-size + - "128" + - --num-hidden-layers + - "6" + - --num-attention-heads + - "4" + - --intermediate-size + - "512" + - --position-embeddings + - learnable + - --lr + - "5e-3" + - --epochs + - "100" + - --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 + - /mnt/d/NCKH/weights/lmc/cifar100 + - --data-path + - /mnt/d/NCKH/datasets/cifar100 + codePath: src/imagenet/train_model.py + codePathLocal: src/imagenet/train_model.py + cpu_count: 24 + cpu_count_logical: 24 + cudaVersion: "13.0" + disk: + /: + total: "1081101176832" + used: "67541254144" + email: khanhvinhbui0512@gmail.com + executable: /home/khanhvinh9/miniconda3/envs/lmc/bin/python + git: + commit: 503490314d9f0bef18a30ec7957b3949bd36c751 + remote: https://github.com/MLResearchX/lmc-transformer.git + gpu: NVIDIA GeForce RTX 5060 Laptop GPU + gpu_count: 1 + gpu_nvidia: + - architecture: Blackwell + memoryTotal: "8546942976" + name: NVIDIA GeForce RTX 5060 Laptop GPU + uuid: GPU-6363d584-7737-137f-cc48-e39ef1a9efbd + host: DESKTOP-UCJ5IKF + memory: + total: "16466997248" + os: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 + program: /home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py + python: CPython 3.10.19 + root: /home/khanhvinh9/lmc-transformer + startedAt: "2025-12-25T05:16:21.910890Z" + writerId: gxr0tuqothrxfogjmk2mbhag5kge133z + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "2": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "3": + - 13 + - 61 + "4": 3.10.19 + "5": 0.23.0 + "6": 4.56.1 + "8": + - 2 + "12": 0.23.0 + "13": linux-x86_64 diff --git a/wandb/run-20251225_121621-23cpikjp/files/diff.patch b/wandb/run-20251225_121621-23cpikjp/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..00762339a061f1e7187065f66f071e30db8800f7 --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/files/diff.patch @@ -0,0 +1,137 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_121621-23cpikjp/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251225_121621-23cpikjp/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..00762339a061f1e7187065f66f071e30db8800f7 --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,137 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_121621-23cpikjp/files/output.log b/wandb/run-20251225_121621-23cpikjp/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..1ab1fe8b885720dd888e8ba2d4323e52f245e2c9 --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/files/output.log @@ -0,0 +1,85 @@ +TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We recommend migrating to PyTorch classes or pinning your version of Transformers. +Starting training... +JAX devices: [CudaDevice(id=0)] +Using 1 devices +Epoch 0 +WARNING:absl:[process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 1 +WARNING:absl:[process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 2 +WARNING:absl:[process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 3 +WARNING:absl:[process=0][thread=MainThread][operation_id=4] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 4 +WARNING:absl:[process=0][thread=MainThread][operation_id=5] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 5 +WARNING:absl:[process=0][thread=MainThread][operation_id=6] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 6 +WARNING:absl:[process=0][thread=MainThread][operation_id=7] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 7 +WARNING:absl:[process=0][thread=MainThread][operation_id=8] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 8 +WARNING:absl:[process=0][thread=MainThread][operation_id=9] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 9 +WARNING:absl:[process=0][thread=MainThread][operation_id=10] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 10 +WARNING:absl:[process=0][thread=MainThread][operation_id=11] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 11 +WARNING:absl:[process=0][thread=MainThread][operation_id=12] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 12 +WARNING:absl:[process=0][thread=MainThread][operation_id=13] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 13 +WARNING:absl:[process=0][thread=MainThread][operation_id=14] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 14 +WARNING:absl:[process=0][thread=MainThread][operation_id=15] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 15 +WARNING:absl:[process=0][thread=MainThread][operation_id=16] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 16 +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 181, in main + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 494, in __iter__ + return self._get_iterator() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 427, in _get_iterator + return _MultiProcessingDataLoaderIter(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1170, in __init__ + w.start() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/process.py", line 121, in start + self._popen = self._Popen(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 224, in _Popen + return _default_context.get_context().Process._Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 288, in _Popen + return Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__ + super().__init__(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__ + self._launch(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 62, in _launch + f.write(fp.getbuffer()) +KeyboardInterrupt +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 181, in main + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 494, in __iter__ + return self._get_iterator() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 427, in _get_iterator + return _MultiProcessingDataLoaderIter(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1170, in __init__ + w.start() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/process.py", line 121, in start + self._popen = self._Popen(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 224, in _Popen + return _default_context.get_context().Process._Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 288, in _Popen + return Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__ + super().__init__(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__ + self._launch(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 62, in _launch + f.write(fp.getbuffer()) +KeyboardInterrupt diff --git a/wandb/run-20251225_121621-23cpikjp/files/requirements.txt b/wandb/run-20251225_121621-23cpikjp/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251225_121621-23cpikjp/files/wandb-metadata.json b/wandb/run-20251225_121621-23cpikjp/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..0f304e1a4d7899db7a97fcfaa27fddb5cc406aab --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/files/wandb-metadata.json @@ -0,0 +1,81 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-25T05:16:21.910890Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "128", + "--num-hidden-layers", + "6", + "--num-attention-heads", + "4", + "--intermediate-size", + "512", + "--position-embeddings", + "learnable", + "--lr", + "5e-3", + "--epochs", + "100", + "--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", + "/mnt/d/NCKH/weights/lmc/cifar100", + "--data-path", + "/mnt/d/NCKH/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67541254144" + } + }, + "memory": { + "total": "16466997248" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "gxr0tuqothrxfogjmk2mbhag5kge133z" +} \ No newline at end of file diff --git a/wandb/run-20251225_121621-23cpikjp/files/wandb-summary.json b/wandb/run-20251225_121621-23cpikjp/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..aee208c592b16a2d0e8214180ec41c0aca31e510 --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/files/wandb-summary.json @@ -0,0 +1 @@ +{"_step":6600,"val/loss":4.53125,"_runtime":1920.723293093,"train/loss":4.53125,"train/acc1":1.5625,"val/acc1":4.6875,"val/acc5":14.84375,"train/acc5":8.59375,"_wandb":{"runtime":1920},"_timestamp":1.7666416135344198e+09,"lr":0.004808558616787195} \ No newline at end of file diff --git a/wandb/run-20251225_121621-23cpikjp/logs/debug-core.log b/wandb/run-20251225_121621-23cpikjp/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..2b5bab16f627768c6af06c50ad1ffcbf47bb7fa7 --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/logs/debug-core.log @@ -0,0 +1,12 @@ +{"time":"2025-12-25T12:16:22.168226593+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmpeee40m2g/port-49556.txt","pid":49556,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T12:16:22.169966099+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":49556} +{"time":"2025-12-25T12:16:22.169897022+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-49556-49708-168527185/socket","Net":"unix"}} +{"time":"2025-12-25T12:16:22.259878451+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T12:16:22.273109662+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"23cpikjp","id":"1(@)"} +{"time":"2025-12-25T12:16:22.801191445+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"23cpikjp","id":"1(@)"} +{"time":"2025-12-25T12:47:05.263734718+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T12:47:05.264056+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T12:47:05.264191197+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T12:47:05.264199267+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T12:47:05.264493172+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-49556-49708-168527185/socket","Net":"unix"}} +{"time":"2025-12-25T12:47:06.640583335+07:00","level":"INFO","msg":"server: parent process exited, terminating service process"} diff --git a/wandb/run-20251225_121621-23cpikjp/logs/debug-internal.log b/wandb/run-20251225_121621-23cpikjp/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..fe3b2bdc6ed12572ade8f56ef0e69e2caf97fcc0 --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/logs/debug-internal.log @@ -0,0 +1,7 @@ +{"time":"2025-12-25T12:16:22.273232958+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T12:16:22.800846713+07:00","level":"INFO","msg":"stream: created new stream","id":"23cpikjp"} +{"time":"2025-12-25T12:16:22.801000549+07:00","level":"INFO","msg":"handler: started","stream_id":"23cpikjp"} +{"time":"2025-12-25T12:16:22.801182448+07:00","level":"INFO","msg":"stream: started","id":"23cpikjp"} +{"time":"2025-12-25T12:16:22.801284094+07:00","level":"INFO","msg":"writer: started","stream_id":"23cpikjp"} +{"time":"2025-12-25T12:16:22.801329534+07:00","level":"INFO","msg":"sender: started","stream_id":"23cpikjp"} +{"time":"2025-12-25T12:47:05.264102984+07:00","level":"INFO","msg":"stream: closing","id":"23cpikjp"} diff --git a/wandb/run-20251225_121621-23cpikjp/logs/debug.log b/wandb/run-20251225_121621-23cpikjp/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..2fa8e83cbbb31f9dad0a311ddb5570a2f24a1109 --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/logs/debug.log @@ -0,0 +1,2818 @@ +2025-12-25 12:16:21,917 INFO MainThread:49556 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 12:16:21,917 INFO MainThread:49556 [wandb_setup.py:_flush():80] Configure stats pid to 49556 +2025-12-25 12:16:21,917 INFO MainThread:49556 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 12:16:21,917 INFO MainThread:49556 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 12:16:21,917 INFO MainThread:49556 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 12:16:21,917 INFO MainThread:49556 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_121621-23cpikjp/logs/debug.log +2025-12-25 12:16:21,917 INFO MainThread:49556 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_121621-23cpikjp/logs/debug-internal.log +2025-12-25 12:16:21,917 INFO MainThread:49556 [wandb_init.py:init():840] calling init triggers +2025-12-25 12:16:21,917 INFO MainThread:49556 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 12:16:21,917 INFO MainThread:49556 [wandb_init.py:init():888] starting backend +2025-12-25 12:16:22,259 INFO MainThread:49556 [wandb_init.py:init():891] sending inform_init request +2025-12-25 12:16:22,266 INFO MainThread:49556 [wandb_init.py:init():899] backend started and connected +2025-12-25 12:16:22,268 INFO MainThread:49556 [wandb_init.py:init():969] updated telemetry +2025-12-25 12:16:22,273 INFO MainThread:49556 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 12:16:23,208 INFO MainThread:49556 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 12:16:23,366 INFO MainThread:49556 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 12:16:23,366 INFO MainThread:49556 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 12:16:23,366 INFO MainThread:49556 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 12:16:23,366 INFO MainThread:49556 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 12:16:23,370 INFO MainThread:49556 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-25 12:47:05,263 INFO wandb-AsyncioManager-main:49556 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-25 12:47:05,263 INFO wandb-AsyncioManager-main:49556 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. +2025-12-25 12:47:05,696 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,702 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,703 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,704 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,704 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,712 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,714 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,714 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,714 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,714 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,714 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,715 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,715 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,715 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,715 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,716 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,716 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,716 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,716 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,716 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,716 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,717 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,717 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,717 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,717 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,717 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,717 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,717 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,718 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,718 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,718 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,718 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,718 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,718 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,719 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,765 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,765 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,766 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,766 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,766 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,766 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,767 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,767 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,767 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,767 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,767 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,767 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,767 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,768 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,768 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,768 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,768 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,768 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,768 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,768 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,768 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,769 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,769 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,769 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,769 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,769 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,770 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,770 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,770 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,770 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,770 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,771 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,771 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,771 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,771 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,771 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,771 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,772 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,772 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,772 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,772 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,772 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,772 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,773 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,773 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,773 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,773 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,773 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,774 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,774 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,781 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,781 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,781 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,781 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,782 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,782 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,782 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,782 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,782 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:05,782 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,044 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,080 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,194 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,194 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,194 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,195 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,195 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,195 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,195 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,195 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,196 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,196 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,196 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,196 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,196 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,197 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,197 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,197 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,197 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,197 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,197 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,198 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,198 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,217 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,218 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,218 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,218 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,218 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,218 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,218 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,219 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,219 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,219 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,219 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,219 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,219 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,219 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,220 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,223 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,223 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,224 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,224 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,224 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,224 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,224 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,225 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,225 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,225 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,225 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,225 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,225 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,225 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,226 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,226 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,226 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,226 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,226 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,226 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,226 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,227 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,227 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,228 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,228 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,228 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,228 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,229 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,229 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,229 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,229 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,229 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,229 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,229 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,230 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,230 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,230 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,230 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,230 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,230 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,230 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,231 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,231 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,231 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,231 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,231 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,231 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,231 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,232 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,232 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,233 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,233 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,233 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,233 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,233 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,233 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,234 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,234 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,234 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,234 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,234 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,235 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,235 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,235 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,236 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,236 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,236 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,236 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,237 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,237 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,237 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,237 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,238 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,238 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,238 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,239 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,239 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,239 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,239 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,239 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,239 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 12:47:06,240 ERROR wandb-AsyncioManager-main:49556 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost diff --git a/wandb/run-20251225_121621-23cpikjp/run-23cpikjp.wandb b/wandb/run-20251225_121621-23cpikjp/run-23cpikjp.wandb new file mode 100644 index 0000000000000000000000000000000000000000..5d8d7c0ae89923ff0d8249fe9c7ee701a573de9a --- /dev/null +++ b/wandb/run-20251225_121621-23cpikjp/run-23cpikjp.wandb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7e498d715d6451e50fc9c0b28a97f65ddd435e586f9d8a601e1352e66ca91833 +size 1376256 diff --git a/wandb/run-20251225_124757-209hc11t/files/code/src/imagenet/train_model.py b/wandb/run-20251225_124757-209hc11t/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c48bfefd7467518c68801e1068bea786a32c1e92 --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/files/code/src/imagenet/train_model.py @@ -0,0 +1,264 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251225_124757-209hc11t/files/config.yaml b/wandb/run-20251225_124757-209hc11t/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..181ee8accc795d9481eb62f6ab6ec6fc43cf41b6 --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/files/config.yaml @@ -0,0 +1,116 @@ +_wandb: + value: + cli_version: 0.23.0 + code_path: code/src/imagenet/train_model.py + e: + oix7h0fjkh4qzybh2hvssvx5ddjg93r1: + args: + - --input-size + - "32" + - --data-set + - CIFAR + - --patch-size + - "4" + - --hidden-size + - "128" + - --num-hidden-layers + - "6" + - --warmup-epochs + - "0" + - --num-attention-heads + - "4" + - --intermediate-size + - "512" + - --position-embeddings + - learnable + - --num-labels + - "100" + - --lr + - "5e-3" + - --epochs + - "100" + - --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 + - /mnt/d/NCKH/weights/lmc/cifar100 + - --data-path + - /mnt/d/NCKH/datasets/cifar100 + codePath: src/imagenet/train_model.py + codePathLocal: src/imagenet/train_model.py + cpu_count: 24 + cpu_count_logical: 24 + cudaVersion: "13.0" + disk: + /: + total: "1081101176832" + used: "67543101440" + email: khanhvinhbui0512@gmail.com + executable: /home/khanhvinh9/miniconda3/envs/lmc/bin/python + git: + commit: 503490314d9f0bef18a30ec7957b3949bd36c751 + remote: https://github.com/MLResearchX/lmc-transformer.git + gpu: NVIDIA GeForce RTX 5060 Laptop GPU + gpu_count: 1 + gpu_nvidia: + - architecture: Blackwell + memoryTotal: "8546942976" + name: NVIDIA GeForce RTX 5060 Laptop GPU + uuid: GPU-6363d584-7737-137f-cc48-e39ef1a9efbd + host: DESKTOP-UCJ5IKF + memory: + total: "16466997248" + os: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 + program: /home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py + python: CPython 3.10.19 + root: /home/khanhvinh9/lmc-transformer + startedAt: "2025-12-25T05:47:57.942574Z" + writerId: oix7h0fjkh4qzybh2hvssvx5ddjg93r1 + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "2": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "3": + - 13 + - 61 + "4": 3.10.19 + "5": 0.23.0 + "6": 4.56.1 + "8": + - 2 + "12": 0.23.0 + "13": linux-x86_64 diff --git a/wandb/run-20251225_124757-209hc11t/files/diff.patch b/wandb/run-20251225_124757-209hc11t/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..00762339a061f1e7187065f66f071e30db8800f7 --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/files/diff.patch @@ -0,0 +1,137 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_124757-209hc11t/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251225_124757-209hc11t/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..00762339a061f1e7187065f66f071e30db8800f7 --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,137 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_124757-209hc11t/files/output.log b/wandb/run-20251225_124757-209hc11t/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..131f794ee955b8625633da818cc1e10e145d36b5 --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/files/output.log @@ -0,0 +1,55 @@ +TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We recommend migrating to PyTorch classes or pinning your version of Transformers. +Starting training... +JAX devices: [CudaDevice(id=0)] +Using 1 devices +Epoch 0 +WARNING:absl:[process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 1 +WARNING:absl:[process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 2 +WARNING:absl:[process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 3 +WARNING:absl:[process=0][thread=MainThread][operation_id=4] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 4 +WARNING:absl:[process=0][thread=MainThread][operation_id=5] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 5 +WARNING:absl:[process=0][thread=MainThread][operation_id=6] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 6 +WARNING:absl:[process=0][thread=MainThread][operation_id=7] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 7 +WARNING:absl:[process=0][thread=MainThread][operation_id=8] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 8 +WARNING:absl:[process=0][thread=MainThread][operation_id=9] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 9 +WARNING:absl:[process=0][thread=MainThread][operation_id=10] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 10 +WARNING:absl:[process=0][thread=MainThread][operation_id=11] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 11 +WARNING:absl:[process=0][thread=MainThread][operation_id=12] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 12 +WARNING:absl:[process=0][thread=MainThread][operation_id=13] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 13 +WARNING:absl:[process=0][thread=MainThread][operation_id=14] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 14 +WARNING:absl:[process=0][thread=MainThread][operation_id=15] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 15 +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 207, in main + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/flax/training/checkpoints.py", line 693, in save_checkpoint + orbax_checkpointer.save( + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 254, in save + raise ValueError(f'Destination {directory} already exists.') +ValueError: Destination /mnt/d/NCKH/weights/lmc/cifar100/lr0.005-learnable-epochs100-batch128-seed0/last_6240 already exists. +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 207, in main + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/flax/training/checkpoints.py", line 693, in save_checkpoint + orbax_checkpointer.save( + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 254, in save + raise ValueError(f'Destination {directory} already exists.') +ValueError: Destination /mnt/d/NCKH/weights/lmc/cifar100/lr0.005-learnable-epochs100-batch128-seed0/last_6240 already exists. diff --git a/wandb/run-20251225_124757-209hc11t/files/requirements.txt b/wandb/run-20251225_124757-209hc11t/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251225_124757-209hc11t/files/wandb-metadata.json b/wandb/run-20251225_124757-209hc11t/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..a6756b78773f2906e3abe20d13573463c30725c9 --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/files/wandb-metadata.json @@ -0,0 +1,85 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-25T05:47:57.942574Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "128", + "--num-hidden-layers", + "6", + "--warmup-epochs", + "0", + "--num-attention-heads", + "4", + "--intermediate-size", + "512", + "--position-embeddings", + "learnable", + "--num-labels", + "100", + "--lr", + "5e-3", + "--epochs", + "100", + "--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", + "/mnt/d/NCKH/weights/lmc/cifar100", + "--data-path", + "/mnt/d/NCKH/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67543101440" + } + }, + "memory": { + "total": "16466997248" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "oix7h0fjkh4qzybh2hvssvx5ddjg93r1" +} \ No newline at end of file diff --git a/wandb/run-20251225_124757-209hc11t/files/wandb-summary.json b/wandb/run-20251225_124757-209hc11t/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..0db06d93f03c69bb72ab0a2d970c83a367e6448e --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/files/wandb-summary.json @@ -0,0 +1 @@ +{"_runtime":1769.471618113,"_step":6240,"train/loss":4.625,"_timestamp":1.7666433772049794e+09,"val/loss":4.59375,"train/acc1":0,"train/acc5":2.34375,"lr":0.004695246461778879,"val/acc1":1.5625,"val/acc5":4.6875,"_wandb":{"runtime":1769}} \ No newline at end of file diff --git a/wandb/run-20251225_124757-209hc11t/logs/debug-core.log b/wandb/run-20251225_124757-209hc11t/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..102505a3132e9f83f56b9d2f75f9f2a7a79e184e --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-25T12:47:58.046395876+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmpf90kaofa/port-96882.txt","pid":96882,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T12:47:58.047586335+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":96882} +{"time":"2025-12-25T12:47:58.04760771+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-96882-96985-2767130955/socket","Net":"unix"}} +{"time":"2025-12-25T12:47:58.227494208+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T12:47:58.237148256+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"209hc11t","id":"1(@)"} +{"time":"2025-12-25T12:47:58.683901993+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"209hc11t","id":"1(@)"} +{"time":"2025-12-25T13:16:17.244487147+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T13:16:17.244570791+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T13:16:17.244555044+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T13:16:17.244666276+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T13:16:17.244674415+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-96882-96985-2767130955/socket","Net":"unix"}} +{"time":"2025-12-25T13:16:19.494390831+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-25T13:16:19.494480884+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-25T13:16:19.494512019+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/run-20251225_124757-209hc11t/logs/debug-internal.log b/wandb/run-20251225_124757-209hc11t/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..952dbb6c7b473a8bb3ffc44d707a73cea9dcf8d4 --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/logs/debug-internal.log @@ -0,0 +1,11 @@ +{"time":"2025-12-25T12:47:58.237324098+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T12:47:58.683609658+07:00","level":"INFO","msg":"stream: created new stream","id":"209hc11t"} +{"time":"2025-12-25T12:47:58.683729171+07:00","level":"INFO","msg":"handler: started","stream_id":"209hc11t"} +{"time":"2025-12-25T12:47:58.683887385+07:00","level":"INFO","msg":"stream: started","id":"209hc11t"} +{"time":"2025-12-25T12:47:58.683980422+07:00","level":"INFO","msg":"writer: started","stream_id":"209hc11t"} +{"time":"2025-12-25T12:47:58.683987835+07:00","level":"INFO","msg":"sender: started","stream_id":"209hc11t"} +{"time":"2025-12-25T13:16:17.244578357+07:00","level":"INFO","msg":"stream: closing","id":"209hc11t"} +{"time":"2025-12-25T13:16:18.786364734+07:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"} +{"time":"2025-12-25T13:16:19.491915045+07:00","level":"INFO","msg":"handler: closed","stream_id":"209hc11t"} +{"time":"2025-12-25T13:16:19.492091422+07:00","level":"INFO","msg":"sender: closed","stream_id":"209hc11t"} +{"time":"2025-12-25T13:16:19.492104905+07:00","level":"INFO","msg":"stream: closed","id":"209hc11t"} diff --git a/wandb/run-20251225_124757-209hc11t/logs/debug.log b/wandb/run-20251225_124757-209hc11t/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..000a46188966d816ed220e58ea424ecf39b67bca --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/logs/debug.log @@ -0,0 +1,23 @@ +2025-12-25 12:47:57,945 INFO MainThread:96882 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 12:47:57,945 INFO MainThread:96882 [wandb_setup.py:_flush():80] Configure stats pid to 96882 +2025-12-25 12:47:57,945 INFO MainThread:96882 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 12:47:57,945 INFO MainThread:96882 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 12:47:57,945 INFO MainThread:96882 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 12:47:57,945 INFO MainThread:96882 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_124757-209hc11t/logs/debug.log +2025-12-25 12:47:57,945 INFO MainThread:96882 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_124757-209hc11t/logs/debug-internal.log +2025-12-25 12:47:57,946 INFO MainThread:96882 [wandb_init.py:init():840] calling init triggers +2025-12-25 12:47:57,946 INFO MainThread:96882 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 12:47:57,946 INFO MainThread:96882 [wandb_init.py:init():888] starting backend +2025-12-25 12:47:58,227 INFO MainThread:96882 [wandb_init.py:init():891] sending inform_init request +2025-12-25 12:47:58,232 INFO MainThread:96882 [wandb_init.py:init():899] backend started and connected +2025-12-25 12:47:58,235 INFO MainThread:96882 [wandb_init.py:init():969] updated telemetry +2025-12-25 12:47:58,243 INFO MainThread:96882 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 12:47:59,124 INFO MainThread:96882 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 12:47:59,232 INFO MainThread:96882 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 12:47:59,232 INFO MainThread:96882 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 12:47:59,232 INFO MainThread:96882 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 12:47:59,232 INFO MainThread:96882 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 12:47:59,235 INFO MainThread:96882 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-25 13:16:17,244 INFO wandb-AsyncioManager-main:96882 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-25 13:16:17,244 INFO wandb-AsyncioManager-main:96882 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. diff --git a/wandb/run-20251225_124757-209hc11t/run-209hc11t.wandb b/wandb/run-20251225_124757-209hc11t/run-209hc11t.wandb new file mode 100644 index 0000000000000000000000000000000000000000..71fe29ccf1e80991a176897d021ad946f4a105a4 --- /dev/null +++ b/wandb/run-20251225_124757-209hc11t/run-209hc11t.wandb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f73303f033dec1958161d8deeeaa1b3512018e99c8674003bf5806373ccec00b +size 1291060 diff --git a/wandb/run-20251225_135106-poxltdws/files/code/src/imagenet/train_model.py b/wandb/run-20251225_135106-poxltdws/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c48bfefd7467518c68801e1068bea786a32c1e92 --- /dev/null +++ b/wandb/run-20251225_135106-poxltdws/files/code/src/imagenet/train_model.py @@ -0,0 +1,264 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251225_135106-poxltdws/files/diff.patch b/wandb/run-20251225_135106-poxltdws/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..00762339a061f1e7187065f66f071e30db8800f7 --- /dev/null +++ b/wandb/run-20251225_135106-poxltdws/files/diff.patch @@ -0,0 +1,137 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_135106-poxltdws/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251225_135106-poxltdws/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..00762339a061f1e7187065f66f071e30db8800f7 --- /dev/null +++ b/wandb/run-20251225_135106-poxltdws/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,137 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_135106-poxltdws/files/output.log b/wandb/run-20251225_135106-poxltdws/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..874297a2da6e15caa16c8a6553fb0c7d624c4357 --- /dev/null +++ b/wandb/run-20251225_135106-poxltdws/files/output.log @@ -0,0 +1,31 @@ +TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We recommend migrating to PyTorch classes or pinning your version of Transformers. +Starting training... +JAX devices: [CudaDevice(id=0)] +Using 1 devices +Epoch 0 +WARNING:absl:[process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 1 +WARNING:absl:[process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 2 +WARNING:absl:[process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 3 +WARNING:absl:[process=0][thread=MainThread][operation_id=4] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 4 +WARNING:absl:[process=0][thread=MainThread][operation_id=5] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 5 +WARNING:absl:[process=0][thread=MainThread][operation_id=6] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 6 +WARNING:absl:[process=0][thread=MainThread][operation_id=7] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 7 +WARNING:absl:[process=0][thread=MainThread][operation_id=8] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 8 +WARNING:absl:[process=0][thread=MainThread][operation_id=9] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 9 +WARNING:absl:[process=0][thread=MainThread][operation_id=10] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 10 +WARNING:absl:[process=0][thread=MainThread][operation_id=11] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 11 +WARNING:absl:[process=0][thread=MainThread][operation_id=12] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 12 +WARNING:absl:[process=0][thread=MainThread][operation_id=13] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 13 diff --git a/wandb/run-20251225_135106-poxltdws/files/requirements.txt b/wandb/run-20251225_135106-poxltdws/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251225_135106-poxltdws/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251225_135106-poxltdws/files/wandb-metadata.json b/wandb/run-20251225_135106-poxltdws/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..78907dbc1b09b033e17968c66417235dbfb3b670 --- /dev/null +++ b/wandb/run-20251225_135106-poxltdws/files/wandb-metadata.json @@ -0,0 +1,85 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-25T06:51:06.044927Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "128", + "--num-hidden-layers", + "6", + "--warmup-epochs", + "0", + "--num-attention-heads", + "4", + "--intermediate-size", + "512", + "--position-embeddings", + "learnable", + "--num-labels", + "100", + "--lr", + "5e-3", + "--epochs", + "100", + "--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", + "/mnt/d/NCKH/weights/lmc/cifar100", + "--data-path", + "/mnt/d/NCKH/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67540455424" + } + }, + "memory": { + "total": "16466997248" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "tw99v9ancb8npzqzrexyr36rhmarycng" +} \ No newline at end of file diff --git a/wandb/run-20251225_135106-poxltdws/logs/debug-core.log b/wandb/run-20251225_135106-poxltdws/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..b94c4248ee2eb57303a3c5c3d9cd1eebbbd0185d --- /dev/null +++ b/wandb/run-20251225_135106-poxltdws/logs/debug-core.log @@ -0,0 +1,6 @@ +{"time":"2025-12-25T13:51:06.314840526+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmpmk25v9e6/port-153485.txt","pid":153485,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T13:51:06.316988397+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":153485} +{"time":"2025-12-25T13:51:06.316987456+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-153485-153640-334279586/socket","Net":"unix"}} +{"time":"2025-12-25T13:51:06.39900732+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T13:51:06.412194147+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"poxltdws","id":"1(@)"} +{"time":"2025-12-25T13:51:06.936334861+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"poxltdws","id":"1(@)"} diff --git a/wandb/run-20251225_135106-poxltdws/logs/debug-internal.log b/wandb/run-20251225_135106-poxltdws/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..87fa1838e0eda81cd6875cd79a71906051b05533 --- /dev/null +++ b/wandb/run-20251225_135106-poxltdws/logs/debug-internal.log @@ -0,0 +1,6 @@ +{"time":"2025-12-25T13:51:06.412339217+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T13:51:06.935972232+07:00","level":"INFO","msg":"stream: created new stream","id":"poxltdws"} +{"time":"2025-12-25T13:51:06.936170216+07:00","level":"INFO","msg":"handler: started","stream_id":"poxltdws"} +{"time":"2025-12-25T13:51:06.936324666+07:00","level":"INFO","msg":"stream: started","id":"poxltdws"} +{"time":"2025-12-25T13:51:06.93642721+07:00","level":"INFO","msg":"writer: started","stream_id":"poxltdws"} +{"time":"2025-12-25T13:51:06.936457681+07:00","level":"INFO","msg":"sender: started","stream_id":"poxltdws"} diff --git a/wandb/run-20251225_135106-poxltdws/logs/debug.log b/wandb/run-20251225_135106-poxltdws/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..a2bbf7848e730d5c815c4cd408a2aaaeb87f042c --- /dev/null +++ b/wandb/run-20251225_135106-poxltdws/logs/debug.log @@ -0,0 +1,21 @@ +2025-12-25 13:51:06,050 INFO MainThread:153485 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 13:51:06,051 INFO MainThread:153485 [wandb_setup.py:_flush():80] Configure stats pid to 153485 +2025-12-25 13:51:06,051 INFO MainThread:153485 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 13:51:06,051 INFO MainThread:153485 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 13:51:06,051 INFO MainThread:153485 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 13:51:06,051 INFO MainThread:153485 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_135106-poxltdws/logs/debug.log +2025-12-25 13:51:06,051 INFO MainThread:153485 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_135106-poxltdws/logs/debug-internal.log +2025-12-25 13:51:06,051 INFO MainThread:153485 [wandb_init.py:init():840] calling init triggers +2025-12-25 13:51:06,051 INFO MainThread:153485 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 13:51:06,051 INFO MainThread:153485 [wandb_init.py:init():888] starting backend +2025-12-25 13:51:06,397 INFO MainThread:153485 [wandb_init.py:init():891] sending inform_init request +2025-12-25 13:51:06,405 INFO MainThread:153485 [wandb_init.py:init():899] backend started and connected +2025-12-25 13:51:06,406 INFO MainThread:153485 [wandb_init.py:init():969] updated telemetry +2025-12-25 13:51:06,412 INFO MainThread:153485 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 13:51:07,496 INFO MainThread:153485 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 13:51:07,654 INFO MainThread:153485 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 13:51:07,654 INFO MainThread:153485 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 13:51:07,654 INFO MainThread:153485 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 13:51:07,654 INFO MainThread:153485 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 13:51:07,658 INFO MainThread:153485 [wandb_init.py:init():1080] run started, returning control to user process diff --git a/wandb/run-20251225_135106-poxltdws/run-poxltdws.wandb b/wandb/run-20251225_135106-poxltdws/run-poxltdws.wandb new file mode 100644 index 0000000000000000000000000000000000000000..7987874fc36a79d22d3dbe45ab0556d5d658c025 --- /dev/null +++ b/wandb/run-20251225_135106-poxltdws/run-poxltdws.wandb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8e3162c9eede16d36f695f7761acaad4ed6f5c44164d771a7d37490e170e0af4 +size 1048576 diff --git a/wandb/run-20251225_144139-zgvpe4a2/files/code/src/imagenet/train_model.py b/wandb/run-20251225_144139-zgvpe4a2/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c48bfefd7467518c68801e1068bea786a32c1e92 --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/files/code/src/imagenet/train_model.py @@ -0,0 +1,264 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251225_144139-zgvpe4a2/files/config.yaml b/wandb/run-20251225_144139-zgvpe4a2/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a71401f34f953f17b5fd2e537f3c5043adb2e990 --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/files/config.yaml @@ -0,0 +1,116 @@ +_wandb: + value: + cli_version: 0.23.0 + code_path: code/src/imagenet/train_model.py + e: + ve5ec04alj8i67pr869vpa5oakm2whrk: + args: + - --input-size + - "32" + - --data-set + - CIFAR + - --patch-size + - "4" + - --hidden-size + - "512" + - --num-hidden-layers + - "6" + - --warmup-epochs + - "5" + - --num-attention-heads + - "8" + - --intermediate-size + - "2048" + - --position-embeddings + - learnable + - --num-labels + - "100" + - --lr + - "1e-4" + - --epochs + - "100" + - --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 + - /mnt/d/NCKH/weights/lmc/cifar100 + - --data-path + - /mnt/d/NCKH/datasets/cifar100 + codePath: src/imagenet/train_model.py + codePathLocal: src/imagenet/train_model.py + cpu_count: 24 + cpu_count_logical: 24 + cudaVersion: "13.0" + disk: + /: + total: "1081101176832" + used: "67536863232" + email: khanhvinhbui0512@gmail.com + executable: /home/khanhvinh9/miniconda3/envs/lmc/bin/python + git: + commit: 503490314d9f0bef18a30ec7957b3949bd36c751 + remote: https://github.com/MLResearchX/lmc-transformer.git + gpu: NVIDIA GeForce RTX 5060 Laptop GPU + gpu_count: 1 + gpu_nvidia: + - architecture: Blackwell + memoryTotal: "8546942976" + name: NVIDIA GeForce RTX 5060 Laptop GPU + uuid: GPU-6363d584-7737-137f-cc48-e39ef1a9efbd + host: DESKTOP-UCJ5IKF + memory: + total: "16466989056" + os: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 + program: /home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py + python: CPython 3.10.19 + root: /home/khanhvinh9/lmc-transformer + startedAt: "2025-12-25T07:41:39.905188Z" + writerId: ve5ec04alj8i67pr869vpa5oakm2whrk + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "2": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "3": + - 13 + - 61 + "4": 3.10.19 + "5": 0.23.0 + "6": 4.56.1 + "8": + - 2 + "12": 0.23.0 + "13": linux-x86_64 diff --git a/wandb/run-20251225_144139-zgvpe4a2/files/diff.patch b/wandb/run-20251225_144139-zgvpe4a2/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..00762339a061f1e7187065f66f071e30db8800f7 --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/files/diff.patch @@ -0,0 +1,137 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_144139-zgvpe4a2/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251225_144139-zgvpe4a2/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..00762339a061f1e7187065f66f071e30db8800f7 --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,137 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..71cd45b 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3861b6b 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_144139-zgvpe4a2/files/output.log b/wandb/run-20251225_144139-zgvpe4a2/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..cd810d6c56ffdd59f2a4f4850b313623986b7c60 --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/files/output.log @@ -0,0 +1,171 @@ +TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We recommend migrating to PyTorch classes or pinning your version of Transformers. +Starting training... +JAX devices: [CudaDevice(id=0)] +Using 1 devices +Epoch 0 +WARNING:absl:[process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 1 +WARNING:absl:[process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 2 +WARNING:absl:[process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 3 +WARNING:absl:[process=0][thread=MainThread][operation_id=4] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 4 +WARNING:absl:[process=0][thread=MainThread][operation_id=5] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 5 +WARNING:absl:[process=0][thread=MainThread][operation_id=6] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 6 +WARNING:absl:[process=0][thread=MainThread][operation_id=7] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 7 +WARNING:absl:[process=0][thread=MainThread][operation_id=8] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 8 +WARNING:absl:[process=0][thread=MainThread][operation_id=9] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 9 +WARNING:absl:[process=0][thread=MainThread][operation_id=10] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 10 +WARNING:absl:[process=0][thread=MainThread][operation_id=11] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 11 +WARNING:absl:[process=0][thread=MainThread][operation_id=12] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 12 +WARNING:absl:[process=0][thread=MainThread][operation_id=13] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 13 +WARNING:absl:[process=0][thread=MainThread][operation_id=14] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 14 +WARNING:absl:[process=0][thread=MainThread][operation_id=15] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 15 +WARNING:absl:[process=0][thread=MainThread][operation_id=16] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 16 +WARNING:absl:[process=0][thread=MainThread][operation_id=17] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 17 +WARNING:absl:[process=0][thread=MainThread][operation_id=18] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 18 +WARNING:absl:[process=0][thread=MainThread][operation_id=19] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 19 +WARNING:absl:[process=0][thread=MainThread][operation_id=20] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 20 +WARNING:absl:[process=0][thread=MainThread][operation_id=21] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 21 +WARNING:absl:[process=0][thread=MainThread][operation_id=22] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 22 +WARNING:absl:[process=0][thread=MainThread][operation_id=23] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 23 +WARNING:absl:[process=0][thread=MainThread][operation_id=24] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 24 +WARNING:absl:[process=0][thread=MainThread][operation_id=25] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 25 +WARNING:absl:[process=0][thread=MainThread][operation_id=26] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 26 +WARNING:absl:[process=0][thread=MainThread][operation_id=27] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 27 +WARNING:absl:[process=0][thread=MainThread][operation_id=28] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 28 +WARNING:absl:[process=0][thread=MainThread][operation_id=29] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 29 +WARNING:absl:[process=0][thread=MainThread][operation_id=30] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 30 +WARNING:absl:[process=0][thread=MainThread][operation_id=31] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 31 +WARNING:absl:[process=0][thread=MainThread][operation_id=32] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 32 +WARNING:absl:[process=0][thread=MainThread][operation_id=33] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 33 +WARNING:absl:[process=0][thread=MainThread][operation_id=34] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 34 +WARNING:absl:[process=0][thread=MainThread][operation_id=35] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 35 +WARNING:absl:[process=0][thread=MainThread][operation_id=36] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 36 +WARNING:absl:[process=0][thread=MainThread][operation_id=37] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 37 +WARNING:absl:[process=0][thread=MainThread][operation_id=38] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 38 +WARNING:absl:[process=0][thread=MainThread][operation_id=39] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 39 +WARNING:absl:[process=0][thread=MainThread][operation_id=40] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 40 +WARNING:absl:[process=0][thread=MainThread][operation_id=41] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 41 +WARNING:absl:[process=0][thread=MainThread][operation_id=42] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 42 +WARNING:absl:[process=0][thread=MainThread][operation_id=43] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 43 +WARNING:absl:[process=0][thread=MainThread][operation_id=44] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 44 +WARNING:absl:[process=0][thread=MainThread][operation_id=45] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 45 +WARNING:absl:[process=0][thread=MainThread][operation_id=46] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 46 +WARNING:absl:[process=0][thread=MainThread][operation_id=47] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 47 +WARNING:absl:[process=0][thread=MainThread][operation_id=48] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 48 +WARNING:absl:[process=0][thread=MainThread][operation_id=49] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 49 +WARNING:absl:[process=0][thread=MainThread][operation_id=50] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 50 +WARNING:absl:[process=0][thread=MainThread][operation_id=51] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 51 +WARNING:absl:[process=0][thread=MainThread][operation_id=52] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 52 +WARNING:absl:[process=0][thread=MainThread][operation_id=53] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 53 +WARNING:absl:[process=0][thread=MainThread][operation_id=54] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 54 +WARNING:absl:[process=0][thread=MainThread][operation_id=55] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 55 +WARNING:absl:[process=0][thread=MainThread][operation_id=56] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 56 +WARNING:absl:[process=0][thread=MainThread][operation_id=57] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 57 +WARNING:absl:[process=0][thread=MainThread][operation_id=58] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 58 +WARNING:absl:[process=0][thread=MainThread][operation_id=59] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 59 +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 166, in main + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 494, in __iter__ + return self._get_iterator() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 427, in _get_iterator + return _MultiProcessingDataLoaderIter(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1170, in __init__ + w.start() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/process.py", line 121, in start + self._popen = self._Popen(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 224, in _Popen + return _default_context.get_context().Process._Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 288, in _Popen + return Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__ + super().__init__(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__ + self._launch(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 62, in _launch + f.write(fp.getbuffer()) +KeyboardInterrupt +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 166, in main + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 494, in __iter__ + return self._get_iterator() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 427, in _get_iterator + return _MultiProcessingDataLoaderIter(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1170, in __init__ + w.start() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/process.py", line 121, in start + self._popen = self._Popen(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 224, in _Popen + return _default_context.get_context().Process._Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 288, in _Popen + return Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__ + super().__init__(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__ + self._launch(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 62, in _launch + f.write(fp.getbuffer()) +KeyboardInterrupt diff --git a/wandb/run-20251225_144139-zgvpe4a2/files/requirements.txt b/wandb/run-20251225_144139-zgvpe4a2/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251225_144139-zgvpe4a2/files/wandb-metadata.json b/wandb/run-20251225_144139-zgvpe4a2/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..bcdfc0e143f2f2cd55a4aa3ae5c8564cda3db10f --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/files/wandb-metadata.json @@ -0,0 +1,85 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-25T07:41:39.905188Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "512", + "--num-hidden-layers", + "6", + "--warmup-epochs", + "5", + "--num-attention-heads", + "8", + "--intermediate-size", + "2048", + "--position-embeddings", + "learnable", + "--num-labels", + "100", + "--lr", + "1e-4", + "--epochs", + "100", + "--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", + "/mnt/d/NCKH/weights/lmc/cifar100", + "--data-path", + "/mnt/d/NCKH/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67536863232" + } + }, + "memory": { + "total": "16466989056" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "ve5ec04alj8i67pr869vpa5oakm2whrk" +} \ No newline at end of file diff --git a/wandb/run-20251225_144139-zgvpe4a2/files/wandb-summary.json b/wandb/run-20251225_144139-zgvpe4a2/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..bdaf71bfb077f3bb2a18354ec92cae8da4c2d61d --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/files/wandb-summary.json @@ -0,0 +1 @@ +{"train/acc5":57.8125,"val/acc5":60.9375,"train/acc1":28.90625,"val/acc1":29.6875,"_wandb":{"runtime":8021},"_timestamp":1.7666566136131847e+09,"val/loss":3.015625,"train/loss":2.890625,"lr":4.543880277196877e-05,"_runtime":8021.092458051,"_step":23010} \ No newline at end of file diff --git a/wandb/run-20251225_144139-zgvpe4a2/logs/debug-core.log b/wandb/run-20251225_144139-zgvpe4a2/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..13b345d5e2558950dcd2f6c7a253f5c2db4fd3c2 --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-25T14:41:40.195103179+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmp3cwo8f2_/port-1717.txt","pid":1717,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T14:41:40.197132088+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":1717} +{"time":"2025-12-25T14:41:40.197140951+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-1717-1878-3467058072/socket","Net":"unix"}} +{"time":"2025-12-25T14:41:40.275636927+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T14:41:40.285333055+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"zgvpe4a2","id":"1(@)"} +{"time":"2025-12-25T14:41:40.725203544+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"zgvpe4a2","id":"1(@)"} +{"time":"2025-12-25T16:57:17.450380968+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T16:57:17.450854552+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T16:57:17.450901733+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T16:57:17.451085878+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T16:57:17.451266829+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-1717-1878-3467058072/socket","Net":"unix"}} +{"time":"2025-12-25T16:57:19.332914837+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-25T16:57:19.332997806+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-25T16:57:19.333015771+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/run-20251225_144139-zgvpe4a2/logs/debug-internal.log b/wandb/run-20251225_144139-zgvpe4a2/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..685f80a3213c52a878a86cbd9ce90c2816d6485e --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/logs/debug-internal.log @@ -0,0 +1,11 @@ +{"time":"2025-12-25T14:41:40.285444331+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T14:41:40.725011803+07:00","level":"INFO","msg":"stream: created new stream","id":"zgvpe4a2"} +{"time":"2025-12-25T14:41:40.725106611+07:00","level":"INFO","msg":"handler: started","stream_id":"zgvpe4a2"} +{"time":"2025-12-25T14:41:40.725195925+07:00","level":"INFO","msg":"stream: started","id":"zgvpe4a2"} +{"time":"2025-12-25T14:41:40.725221382+07:00","level":"INFO","msg":"writer: started","stream_id":"zgvpe4a2"} +{"time":"2025-12-25T14:41:40.725246559+07:00","level":"INFO","msg":"sender: started","stream_id":"zgvpe4a2"} +{"time":"2025-12-25T16:57:17.450978588+07:00","level":"INFO","msg":"stream: closing","id":"zgvpe4a2"} +{"time":"2025-12-25T16:57:18.933775392+07:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"} +{"time":"2025-12-25T16:57:19.326219068+07:00","level":"INFO","msg":"handler: closed","stream_id":"zgvpe4a2"} +{"time":"2025-12-25T16:57:19.326351214+07:00","level":"INFO","msg":"sender: closed","stream_id":"zgvpe4a2"} +{"time":"2025-12-25T16:57:19.326390361+07:00","level":"INFO","msg":"stream: closed","id":"zgvpe4a2"} diff --git a/wandb/run-20251225_144139-zgvpe4a2/logs/debug.log b/wandb/run-20251225_144139-zgvpe4a2/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..146a7f5119bdca982658cb91bc2333de48e720ab --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/logs/debug.log @@ -0,0 +1,1115 @@ +2025-12-25 14:41:39,907 INFO MainThread:1717 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 14:41:39,907 INFO MainThread:1717 [wandb_setup.py:_flush():80] Configure stats pid to 1717 +2025-12-25 14:41:39,907 INFO MainThread:1717 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 14:41:39,907 INFO MainThread:1717 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 14:41:39,907 INFO MainThread:1717 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 14:41:39,907 INFO MainThread:1717 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_144139-zgvpe4a2/logs/debug.log +2025-12-25 14:41:39,907 INFO MainThread:1717 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_144139-zgvpe4a2/logs/debug-internal.log +2025-12-25 14:41:39,907 INFO MainThread:1717 [wandb_init.py:init():840] calling init triggers +2025-12-25 14:41:39,907 INFO MainThread:1717 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 14:41:39,907 INFO MainThread:1717 [wandb_init.py:init():888] starting backend +2025-12-25 14:41:40,275 INFO MainThread:1717 [wandb_init.py:init():891] sending inform_init request +2025-12-25 14:41:40,281 INFO MainThread:1717 [wandb_init.py:init():899] backend started and connected +2025-12-25 14:41:40,282 INFO MainThread:1717 [wandb_init.py:init():969] updated telemetry +2025-12-25 14:41:40,286 INFO MainThread:1717 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 14:41:41,120 INFO MainThread:1717 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 14:41:41,269 INFO MainThread:1717 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 14:41:41,270 INFO MainThread:1717 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 14:41:41,270 INFO MainThread:1717 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 14:41:41,270 INFO MainThread:1717 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 14:41:41,273 INFO MainThread:1717 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-25 16:57:17,450 INFO wandb-AsyncioManager-main:1717 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-25 16:57:17,450 INFO wandb-AsyncioManager-main:1717 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. +2025-12-25 16:57:17,792 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:17,797 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,295 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,303 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,303 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,303 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,303 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,304 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,304 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,304 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,304 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,304 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,304 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,305 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,305 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,305 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,305 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,305 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,305 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,305 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,305 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,306 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,306 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,306 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,306 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,306 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,318 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,318 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,318 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,318 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,318 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,319 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,319 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,320 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,320 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,320 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,320 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,320 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,320 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,321 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,321 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,321 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,321 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,321 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,322 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,322 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,322 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,322 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,322 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,322 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,322 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,323 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,323 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,323 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,323 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,323 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,323 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,323 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,324 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,324 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,324 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,324 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,324 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,324 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,325 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,325 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,325 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,325 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,325 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,325 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,325 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,326 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,326 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,326 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,326 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,326 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,326 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,662 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,662 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,663 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,663 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,663 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,663 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-25 16:57:18,664 ERROR wandb-AsyncioManager-main:1717 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost diff --git a/wandb/run-20251225_144139-zgvpe4a2/run-zgvpe4a2.wandb b/wandb/run-20251225_144139-zgvpe4a2/run-zgvpe4a2.wandb new file mode 100644 index 0000000000000000000000000000000000000000..462262c052059c4fd7a7116cdb5a6795598f0f0d --- /dev/null +++ b/wandb/run-20251225_144139-zgvpe4a2/run-zgvpe4a2.wandb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9885a97ffc479c5b0677d7a87c5cceaa527db3a581c2ba621772e3d32901505c +size 6017199 diff --git a/wandb/run-20251225_170452-24ta21sy/files/code/src/imagenet/train_model.py b/wandb/run-20251225_170452-24ta21sy/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c48bfefd7467518c68801e1068bea786a32c1e92 --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/files/code/src/imagenet/train_model.py @@ -0,0 +1,264 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251225_170452-24ta21sy/files/config.yaml b/wandb/run-20251225_170452-24ta21sy/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2183476c80774700c66b85482cac4c6330ae9c37 --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/files/config.yaml @@ -0,0 +1,116 @@ +_wandb: + value: + cli_version: 0.23.0 + code_path: code/src/imagenet/train_model.py + e: + oaxikgbq9svd3895s6nt80ic40i9h3sp: + args: + - --input-size + - "32" + - --data-set + - CIFAR + - --patch-size + - "4" + - --hidden-size + - "512" + - --num-hidden-layers + - "6" + - --warmup-epochs + - "5" + - --num-attention-heads + - "8" + - --intermediate-size + - "2048" + - --position-embeddings + - learnable + - --num-labels + - "100" + - --lr + - "1e-4" + - --epochs + - "100" + - --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 + - /mnt/d/NCKH/weights/lmc/cifar100 + - --data-path + - /mnt/d/NCKH/datasets/cifar100 + codePath: src/imagenet/train_model.py + codePathLocal: src/imagenet/train_model.py + cpu_count: 24 + cpu_count_logical: 24 + cudaVersion: "13.0" + disk: + /: + total: "1081101176832" + used: "67543863296" + email: khanhvinhbui0512@gmail.com + executable: /home/khanhvinh9/miniconda3/envs/lmc/bin/python + git: + commit: 503490314d9f0bef18a30ec7957b3949bd36c751 + remote: https://github.com/MLResearchX/lmc-transformer.git + gpu: NVIDIA GeForce RTX 5060 Laptop GPU + gpu_count: 1 + gpu_nvidia: + - architecture: Blackwell + memoryTotal: "8546942976" + name: NVIDIA GeForce RTX 5060 Laptop GPU + uuid: GPU-6363d584-7737-137f-cc48-e39ef1a9efbd + host: DESKTOP-UCJ5IKF + memory: + total: "16466989056" + os: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 + program: /home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py + python: CPython 3.10.19 + root: /home/khanhvinh9/lmc-transformer + startedAt: "2025-12-25T10:04:52.107174Z" + writerId: oaxikgbq9svd3895s6nt80ic40i9h3sp + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "2": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "3": + - 13 + - 61 + "4": 3.10.19 + "5": 0.23.0 + "6": 4.56.1 + "8": + - 2 + "12": 0.23.0 + "13": linux-x86_64 diff --git a/wandb/run-20251225_170452-24ta21sy/files/diff.patch b/wandb/run-20251225_170452-24ta21sy/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..64375f8ecec60cacb0096f58826df1afb889a5ae --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/files/diff.patch @@ -0,0 +1,214 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..ffbb4f8 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3f8f560 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +@@ -78,33 +79,48 @@ def build_dataset(is_train, args): + + def build_transform(is_train, args): + resize_im = args.input_size > 32 +- if is_train: +- # this should always dispatch to transforms_imagenet_train +- transform = create_transform( +- input_size=args.input_size, +- is_training=True, +- color_jitter=args.color_jitter, +- auto_augment=args.aa, +- interpolation=args.train_interpolation, +- re_prob=args.reprob, +- re_mode=args.remode, +- re_count=args.recount, +- ) +- if not resize_im: +- # replace RandomResizedCropAndInterpolation with +- # RandomCrop +- transform.transforms[0] = transforms.RandomCrop( +- args.input_size, padding=4) +- return transform +- t = [] +- if resize_im: +- size = int((256 / 224) * args.input_size) +- t.append( +- transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images +- ) +- t.append(transforms.CenterCrop(args.input_size)) +- t.append(transforms.ToTensor()) +- t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) +- return transforms.Compose(t) ++ if args.data_set == 'CIFAR': ++ if is_train: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.RandomHorizontalFlip(), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ if is_train: ++ # this should always dispatch to transforms_imagenet_train ++ transform = create_transform( ++ input_size=args.input_size, ++ is_training=True, ++ color_jitter=args.color_jitter, ++ auto_augment=args.aa, ++ interpolation=args.train_interpolation, ++ re_prob=args.reprob, ++ re_mode=args.remode, ++ re_count=args.recount, ++ ) ++ if not resize_im: ++ # replace RandomResizedCropAndInterpolation with ++ # RandomCrop ++ transform.transforms[0] = transforms.RandomCrop( ++ args.input_size, padding=4) ++ return transform ++ t = [] ++ if resize_im: ++ size = int((256 / 224) * args.input_size) ++ t.append( ++ transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images ++ ) ++ t.append(transforms.CenterCrop(args.input_size)) ++ t.append(transforms.ToTensor()) ++ t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) ++ return transforms.Compose(t) + + +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_170452-24ta21sy/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251225_170452-24ta21sy/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..64375f8ecec60cacb0096f58826df1afb889a5ae --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,214 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..ffbb4f8 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3f8f560 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +@@ -78,33 +79,48 @@ def build_dataset(is_train, args): + + def build_transform(is_train, args): + resize_im = args.input_size > 32 +- if is_train: +- # this should always dispatch to transforms_imagenet_train +- transform = create_transform( +- input_size=args.input_size, +- is_training=True, +- color_jitter=args.color_jitter, +- auto_augment=args.aa, +- interpolation=args.train_interpolation, +- re_prob=args.reprob, +- re_mode=args.remode, +- re_count=args.recount, +- ) +- if not resize_im: +- # replace RandomResizedCropAndInterpolation with +- # RandomCrop +- transform.transforms[0] = transforms.RandomCrop( +- args.input_size, padding=4) +- return transform +- t = [] +- if resize_im: +- size = int((256 / 224) * args.input_size) +- t.append( +- transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images +- ) +- t.append(transforms.CenterCrop(args.input_size)) +- t.append(transforms.ToTensor()) +- t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) +- return transforms.Compose(t) ++ if args.data_set == 'CIFAR': ++ if is_train: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.RandomHorizontalFlip(), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ if is_train: ++ # this should always dispatch to transforms_imagenet_train ++ transform = create_transform( ++ input_size=args.input_size, ++ is_training=True, ++ color_jitter=args.color_jitter, ++ auto_augment=args.aa, ++ interpolation=args.train_interpolation, ++ re_prob=args.reprob, ++ re_mode=args.remode, ++ re_count=args.recount, ++ ) ++ if not resize_im: ++ # replace RandomResizedCropAndInterpolation with ++ # RandomCrop ++ transform.transforms[0] = transforms.RandomCrop( ++ args.input_size, padding=4) ++ return transform ++ t = [] ++ if resize_im: ++ size = int((256 / 224) * args.input_size) ++ t.append( ++ transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images ++ ) ++ t.append(transforms.CenterCrop(args.input_size)) ++ t.append(transforms.ToTensor()) ++ t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) ++ return transforms.Compose(t) + + +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_170452-24ta21sy/files/output.log b/wandb/run-20251225_170452-24ta21sy/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..066051227d94928acc5d22f066c1e908155f5e6b --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/files/output.log @@ -0,0 +1,141 @@ +TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We recommend migrating to PyTorch classes or pinning your version of Transformers. +Starting training... +JAX devices: [CudaDevice(id=0)] +Using 1 devices +Epoch 0 +WARNING:absl:[process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 1 +WARNING:absl:[process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 2 +WARNING:absl:[process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 3 +WARNING:absl:[process=0][thread=MainThread][operation_id=4] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 4 +WARNING:absl:[process=0][thread=MainThread][operation_id=5] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 5 +WARNING:absl:[process=0][thread=MainThread][operation_id=6] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 6 +WARNING:absl:[process=0][thread=MainThread][operation_id=7] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 7 +WARNING:absl:[process=0][thread=MainThread][operation_id=8] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 8 +WARNING:absl:[process=0][thread=MainThread][operation_id=9] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 9 +WARNING:absl:[process=0][thread=MainThread][operation_id=10] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 10 +WARNING:absl:[process=0][thread=MainThread][operation_id=11] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 11 +WARNING:absl:[process=0][thread=MainThread][operation_id=12] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 12 +WARNING:absl:[process=0][thread=MainThread][operation_id=13] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 13 +WARNING:absl:[process=0][thread=MainThread][operation_id=14] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 14 +WARNING:absl:[process=0][thread=MainThread][operation_id=15] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 15 +WARNING:absl:[process=0][thread=MainThread][operation_id=16] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 16 +WARNING:absl:[process=0][thread=MainThread][operation_id=17] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 17 +WARNING:absl:[process=0][thread=MainThread][operation_id=18] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 18 +WARNING:absl:[process=0][thread=MainThread][operation_id=19] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 19 +WARNING:absl:[process=0][thread=MainThread][operation_id=20] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 20 +WARNING:absl:[process=0][thread=MainThread][operation_id=21] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 21 +WARNING:absl:[process=0][thread=MainThread][operation_id=22] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 22 +WARNING:absl:[process=0][thread=MainThread][operation_id=23] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 23 +WARNING:absl:[process=0][thread=MainThread][operation_id=24] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 24 +WARNING:absl:[process=0][thread=MainThread][operation_id=25] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 25 +WARNING:absl:[process=0][thread=MainThread][operation_id=26] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 26 +WARNING:absl:[process=0][thread=MainThread][operation_id=27] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 27 +WARNING:absl:[process=0][thread=MainThread][operation_id=28] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 28 +WARNING:absl:[process=0][thread=MainThread][operation_id=29] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 29 +WARNING:absl:[process=0][thread=MainThread][operation_id=30] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 30 +WARNING:absl:[process=0][thread=MainThread][operation_id=31] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 31 +WARNING:absl:[process=0][thread=MainThread][operation_id=32] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 32 +WARNING:absl:[process=0][thread=MainThread][operation_id=33] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 33 +WARNING:absl:[process=0][thread=MainThread][operation_id=34] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 34 +WARNING:absl:[process=0][thread=MainThread][operation_id=35] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 35 +WARNING:absl:[process=0][thread=MainThread][operation_id=36] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 36 +WARNING:absl:[process=0][thread=MainThread][operation_id=37] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 37 +WARNING:absl:[process=0][thread=MainThread][operation_id=38] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 38 +WARNING:absl:[process=0][thread=MainThread][operation_id=39] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 39 +WARNING:absl:[process=0][thread=MainThread][operation_id=40] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 40 +WARNING:absl:[process=0][thread=MainThread][operation_id=41] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 41 +WARNING:absl:[process=0][thread=MainThread][operation_id=42] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 42 +WARNING:absl:[process=0][thread=MainThread][operation_id=43] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 43 +WARNING:absl:[process=0][thread=MainThread][operation_id=44] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 44 +WARNING:absl:[process=0][thread=MainThread][operation_id=45] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 45 +WARNING:absl:[process=0][thread=MainThread][operation_id=46] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 46 +WARNING:absl:[process=0][thread=MainThread][operation_id=47] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 47 +WARNING:absl:[process=0][thread=MainThread][operation_id=48] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 48 +WARNING:absl:[process=0][thread=MainThread][operation_id=49] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 49 +WARNING:absl:[process=0][thread=MainThread][operation_id=50] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 50 +WARNING:absl:[process=0][thread=MainThread][operation_id=51] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 51 +WARNING:absl:[process=0][thread=MainThread][operation_id=52] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 52 +WARNING:absl:[process=0][thread=MainThread][operation_id=53] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 53 +WARNING:absl:[process=0][thread=MainThread][operation_id=54] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 54 +WARNING:absl:[process=0][thread=MainThread][operation_id=55] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 55 +WARNING:absl:[process=0][thread=MainThread][operation_id=56] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 56 +WARNING:absl:[process=0][thread=MainThread][operation_id=57] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 57 +WARNING:absl:[process=0][thread=MainThread][operation_id=58] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 58 +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 207, in main + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/flax/training/checkpoints.py", line 693, in save_checkpoint + orbax_checkpointer.save( + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 254, in save + raise ValueError(f'Destination {directory} already exists.') +ValueError: Destination /mnt/d/NCKH/weights/lmc/cifar100/lr0.0001-learnable-epochs100-batch128-seed0/last_23010 already exists. +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 207, in main + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/flax/training/checkpoints.py", line 693, in save_checkpoint + orbax_checkpointer.save( + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 254, in save + raise ValueError(f'Destination {directory} already exists.') +ValueError: Destination /mnt/d/NCKH/weights/lmc/cifar100/lr0.0001-learnable-epochs100-batch128-seed0/last_23010 already exists. diff --git a/wandb/run-20251225_170452-24ta21sy/files/requirements.txt b/wandb/run-20251225_170452-24ta21sy/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251225_170452-24ta21sy/files/wandb-metadata.json b/wandb/run-20251225_170452-24ta21sy/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..926f16bfd1312b060a88e4ca3b2dd25b061787c0 --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/files/wandb-metadata.json @@ -0,0 +1,85 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-25T10:04:52.107174Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "512", + "--num-hidden-layers", + "6", + "--warmup-epochs", + "5", + "--num-attention-heads", + "8", + "--intermediate-size", + "2048", + "--position-embeddings", + "learnable", + "--num-labels", + "100", + "--lr", + "1e-4", + "--epochs", + "100", + "--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", + "/mnt/d/NCKH/weights/lmc/cifar100", + "--data-path", + "/mnt/d/NCKH/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67543863296" + } + }, + "memory": { + "total": "16466989056" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "oaxikgbq9svd3895s6nt80ic40i9h3sp" +} \ No newline at end of file diff --git a/wandb/run-20251225_170452-24ta21sy/files/wandb-summary.json b/wandb/run-20251225_170452-24ta21sy/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..1e23f1ff8e1f172ae59278a514483800e9bb1ed4 --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/files/wandb-summary.json @@ -0,0 +1 @@ +{"train/loss":0.5859375,"val/acc1":53.90625,"val/loss":1.9140625,"lr":4.543880277196877e-05,"_wandb":{"runtime":7310},"_timestamp":1.7666644853279705e+09,"_step":23010,"train/acc1":82.03125,"train/acc5":95.3125,"val/acc5":79.6875,"_runtime":7310.655955771} \ No newline at end of file diff --git a/wandb/run-20251225_170452-24ta21sy/logs/debug-core.log b/wandb/run-20251225_170452-24ta21sy/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..2f2dcf81ebe49f1830b7d40a7a0ab2aee3edb608 --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-25T17:04:52.220072606+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmpk3cw1om3/port-194427.txt","pid":194427,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T17:04:52.221380016+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":194427} +{"time":"2025-12-25T17:04:52.221348688+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-194427-194566-2564616621/socket","Net":"unix"}} +{"time":"2025-12-25T17:04:52.393554452+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T17:04:52.400641026+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"24ta21sy","id":"1(@)"} +{"time":"2025-12-25T17:04:52.85214769+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"24ta21sy","id":"1(@)"} +{"time":"2025-12-25T19:08:05.371017809+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T19:08:05.371213813+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T19:08:05.371170324+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T19:08:05.371355188+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-194427-194566-2564616621/socket","Net":"unix"}} +{"time":"2025-12-25T19:08:05.371403925+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T19:08:07.094147286+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-25T19:08:07.095253314+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-25T19:08:07.095678982+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/run-20251225_170452-24ta21sy/logs/debug-internal.log b/wandb/run-20251225_170452-24ta21sy/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..e1fe10c4aa006b6dc8b7309a157aac105f708981 --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/logs/debug-internal.log @@ -0,0 +1,11 @@ +{"time":"2025-12-25T17:04:52.400810463+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T17:04:52.851836946+07:00","level":"INFO","msg":"stream: created new stream","id":"24ta21sy"} +{"time":"2025-12-25T17:04:52.852056823+07:00","level":"INFO","msg":"handler: started","stream_id":"24ta21sy"} +{"time":"2025-12-25T17:04:52.852135199+07:00","level":"INFO","msg":"stream: started","id":"24ta21sy"} +{"time":"2025-12-25T17:04:52.8521831+07:00","level":"INFO","msg":"writer: started","stream_id":"24ta21sy"} +{"time":"2025-12-25T17:04:52.852220268+07:00","level":"INFO","msg":"sender: started","stream_id":"24ta21sy"} +{"time":"2025-12-25T19:08:05.37119723+07:00","level":"INFO","msg":"stream: closing","id":"24ta21sy"} +{"time":"2025-12-25T19:08:06.71181412+07:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"} +{"time":"2025-12-25T19:08:07.089402958+07:00","level":"INFO","msg":"handler: closed","stream_id":"24ta21sy"} +{"time":"2025-12-25T19:08:07.089720552+07:00","level":"INFO","msg":"sender: closed","stream_id":"24ta21sy"} +{"time":"2025-12-25T19:08:07.089865778+07:00","level":"INFO","msg":"stream: closed","id":"24ta21sy"} diff --git a/wandb/run-20251225_170452-24ta21sy/logs/debug.log b/wandb/run-20251225_170452-24ta21sy/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..e8f3adddd4f5c1fb0c6eb46612bebbc15eb093a1 --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/logs/debug.log @@ -0,0 +1,23 @@ +2025-12-25 17:04:52,111 INFO MainThread:194427 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 17:04:52,111 INFO MainThread:194427 [wandb_setup.py:_flush():80] Configure stats pid to 194427 +2025-12-25 17:04:52,111 INFO MainThread:194427 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 17:04:52,111 INFO MainThread:194427 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 17:04:52,111 INFO MainThread:194427 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 17:04:52,111 INFO MainThread:194427 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_170452-24ta21sy/logs/debug.log +2025-12-25 17:04:52,111 INFO MainThread:194427 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_170452-24ta21sy/logs/debug-internal.log +2025-12-25 17:04:52,111 INFO MainThread:194427 [wandb_init.py:init():840] calling init triggers +2025-12-25 17:04:52,111 INFO MainThread:194427 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 17:04:52,111 INFO MainThread:194427 [wandb_init.py:init():888] starting backend +2025-12-25 17:04:52,393 INFO MainThread:194427 [wandb_init.py:init():891] sending inform_init request +2025-12-25 17:04:52,398 INFO MainThread:194427 [wandb_init.py:init():899] backend started and connected +2025-12-25 17:04:52,399 INFO MainThread:194427 [wandb_init.py:init():969] updated telemetry +2025-12-25 17:04:52,403 INFO MainThread:194427 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 17:04:53,267 INFO MainThread:194427 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 17:04:53,374 INFO MainThread:194427 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 17:04:53,374 INFO MainThread:194427 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 17:04:53,374 INFO MainThread:194427 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 17:04:53,374 INFO MainThread:194427 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 17:04:53,376 INFO MainThread:194427 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-25 19:08:05,370 INFO wandb-AsyncioManager-main:194427 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-25 19:08:05,370 INFO wandb-AsyncioManager-main:194427 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. diff --git a/wandb/run-20251225_170452-24ta21sy/run-24ta21sy.wandb b/wandb/run-20251225_170452-24ta21sy/run-24ta21sy.wandb new file mode 100644 index 0000000000000000000000000000000000000000..436577720828ec7fceac5a60122ac433ee170e57 --- /dev/null +++ b/wandb/run-20251225_170452-24ta21sy/run-24ta21sy.wandb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ebc26f0c7cfc79ca02999e4585f49c56b8125e227fcb1afafba76142037bda3d +size 6001525 diff --git a/wandb/run-20251225_200829-ad33gu8w/files/code/src/imagenet/train_model.py b/wandb/run-20251225_200829-ad33gu8w/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c48bfefd7467518c68801e1068bea786a32c1e92 --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/files/code/src/imagenet/train_model.py @@ -0,0 +1,264 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251225_200829-ad33gu8w/files/config.yaml b/wandb/run-20251225_200829-ad33gu8w/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..91f85c7f1d0b110aa3253cbfc2c2fc001cfd43a5 --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/files/config.yaml @@ -0,0 +1,116 @@ +_wandb: + value: + cli_version: 0.23.0 + code_path: code/src/imagenet/train_model.py + e: + hxpb1vakuyd64uxdubge65x4g39p6ahi: + args: + - --input-size + - "32" + - --data-set + - CIFAR + - --patch-size + - "4" + - --hidden-size + - "512" + - --num-hidden-layers + - "6" + - --warmup-epochs + - "5" + - --num-attention-heads + - "8" + - --intermediate-size + - "2048" + - --position-embeddings + - learnable + - --num-labels + - "100" + - --lr + - "1e-4" + - --epochs + - "100" + - --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 + - /mnt/d/NCKH/weights/lmc/cifar100 + - --data-path + - /mnt/d/NCKH/datasets/cifar100 + codePath: src/imagenet/train_model.py + codePathLocal: src/imagenet/train_model.py + cpu_count: 24 + cpu_count_logical: 24 + cudaVersion: "13.0" + disk: + /: + total: "1081101176832" + used: "67550371840" + email: khanhvinhbui0512@gmail.com + executable: /home/khanhvinh9/miniconda3/envs/lmc/bin/python + git: + commit: 503490314d9f0bef18a30ec7957b3949bd36c751 + remote: https://github.com/MLResearchX/lmc-transformer.git + gpu: NVIDIA GeForce RTX 5060 Laptop GPU + gpu_count: 1 + gpu_nvidia: + - architecture: Blackwell + memoryTotal: "8546942976" + name: NVIDIA GeForce RTX 5060 Laptop GPU + uuid: GPU-6363d584-7737-137f-cc48-e39ef1a9efbd + host: DESKTOP-UCJ5IKF + memory: + total: "16466989056" + os: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 + program: /home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py + python: CPython 3.10.19 + root: /home/khanhvinh9/lmc-transformer + startedAt: "2025-12-25T13:08:29.527349Z" + writerId: hxpb1vakuyd64uxdubge65x4g39p6ahi + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "2": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "3": + - 13 + - 61 + "4": 3.10.19 + "5": 0.23.0 + "6": 4.56.1 + "8": + - 2 + "12": 0.23.0 + "13": linux-x86_64 diff --git a/wandb/run-20251225_200829-ad33gu8w/files/diff.patch b/wandb/run-20251225_200829-ad33gu8w/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..64375f8ecec60cacb0096f58826df1afb889a5ae --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/files/diff.patch @@ -0,0 +1,214 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..ffbb4f8 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3f8f560 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +@@ -78,33 +79,48 @@ def build_dataset(is_train, args): + + def build_transform(is_train, args): + resize_im = args.input_size > 32 +- if is_train: +- # this should always dispatch to transforms_imagenet_train +- transform = create_transform( +- input_size=args.input_size, +- is_training=True, +- color_jitter=args.color_jitter, +- auto_augment=args.aa, +- interpolation=args.train_interpolation, +- re_prob=args.reprob, +- re_mode=args.remode, +- re_count=args.recount, +- ) +- if not resize_im: +- # replace RandomResizedCropAndInterpolation with +- # RandomCrop +- transform.transforms[0] = transforms.RandomCrop( +- args.input_size, padding=4) +- return transform +- t = [] +- if resize_im: +- size = int((256 / 224) * args.input_size) +- t.append( +- transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images +- ) +- t.append(transforms.CenterCrop(args.input_size)) +- t.append(transforms.ToTensor()) +- t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) +- return transforms.Compose(t) ++ if args.data_set == 'CIFAR': ++ if is_train: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.RandomHorizontalFlip(), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ if is_train: ++ # this should always dispatch to transforms_imagenet_train ++ transform = create_transform( ++ input_size=args.input_size, ++ is_training=True, ++ color_jitter=args.color_jitter, ++ auto_augment=args.aa, ++ interpolation=args.train_interpolation, ++ re_prob=args.reprob, ++ re_mode=args.remode, ++ re_count=args.recount, ++ ) ++ if not resize_im: ++ # replace RandomResizedCropAndInterpolation with ++ # RandomCrop ++ transform.transforms[0] = transforms.RandomCrop( ++ args.input_size, padding=4) ++ return transform ++ t = [] ++ if resize_im: ++ size = int((256 / 224) * args.input_size) ++ t.append( ++ transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images ++ ) ++ t.append(transforms.CenterCrop(args.input_size)) ++ t.append(transforms.ToTensor()) ++ t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) ++ return transforms.Compose(t) + + +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_200829-ad33gu8w/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251225_200829-ad33gu8w/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..64375f8ecec60cacb0096f58826df1afb889a5ae --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,214 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..ffbb4f8 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3f8f560 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +@@ -78,33 +79,48 @@ def build_dataset(is_train, args): + + def build_transform(is_train, args): + resize_im = args.input_size > 32 +- if is_train: +- # this should always dispatch to transforms_imagenet_train +- transform = create_transform( +- input_size=args.input_size, +- is_training=True, +- color_jitter=args.color_jitter, +- auto_augment=args.aa, +- interpolation=args.train_interpolation, +- re_prob=args.reprob, +- re_mode=args.remode, +- re_count=args.recount, +- ) +- if not resize_im: +- # replace RandomResizedCropAndInterpolation with +- # RandomCrop +- transform.transforms[0] = transforms.RandomCrop( +- args.input_size, padding=4) +- return transform +- t = [] +- if resize_im: +- size = int((256 / 224) * args.input_size) +- t.append( +- transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images +- ) +- t.append(transforms.CenterCrop(args.input_size)) +- t.append(transforms.ToTensor()) +- t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) +- return transforms.Compose(t) ++ if args.data_set == 'CIFAR': ++ if is_train: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.RandomHorizontalFlip(), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ if is_train: ++ # this should always dispatch to transforms_imagenet_train ++ transform = create_transform( ++ input_size=args.input_size, ++ is_training=True, ++ color_jitter=args.color_jitter, ++ auto_augment=args.aa, ++ interpolation=args.train_interpolation, ++ re_prob=args.reprob, ++ re_mode=args.remode, ++ re_count=args.recount, ++ ) ++ if not resize_im: ++ # replace RandomResizedCropAndInterpolation with ++ # RandomCrop ++ transform.transforms[0] = transforms.RandomCrop( ++ args.input_size, padding=4) ++ return transform ++ t = [] ++ if resize_im: ++ size = int((256 / 224) * args.input_size) ++ t.append( ++ transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images ++ ) ++ t.append(transforms.CenterCrop(args.input_size)) ++ t.append(transforms.ToTensor()) ++ t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) ++ return transforms.Compose(t) + + +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_200829-ad33gu8w/files/output.log b/wandb/run-20251225_200829-ad33gu8w/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..fcdeae4681efe5958fff5c066222be1ee18b1fb1 --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/files/output.log @@ -0,0 +1,204 @@ +TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We recommend migrating to PyTorch classes or pinning your version of Transformers. +Starting training... +JAX devices: [CudaDevice(id=0)] +Using 1 devices +Epoch 0 +WARNING:absl:[process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 1 +WARNING:absl:[process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 2 +WARNING:absl:[process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 3 +WARNING:absl:[process=0][thread=MainThread][operation_id=4] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 4 +WARNING:absl:[process=0][thread=MainThread][operation_id=5] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 5 +WARNING:absl:[process=0][thread=MainThread][operation_id=6] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 6 +WARNING:absl:[process=0][thread=MainThread][operation_id=7] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 7 +WARNING:absl:[process=0][thread=MainThread][operation_id=8] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 8 +WARNING:absl:[process=0][thread=MainThread][operation_id=9] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 9 +WARNING:absl:[process=0][thread=MainThread][operation_id=10] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 10 +WARNING:absl:[process=0][thread=MainThread][operation_id=11] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 11 +WARNING:absl:[process=0][thread=MainThread][operation_id=12] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 12 +WARNING:absl:[process=0][thread=MainThread][operation_id=13] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 13 +WARNING:absl:[process=0][thread=MainThread][operation_id=14] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 14 +WARNING:absl:[process=0][thread=MainThread][operation_id=15] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 15 +WARNING:absl:[process=0][thread=MainThread][operation_id=16] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 16 +WARNING:absl:[process=0][thread=MainThread][operation_id=17] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 17 +WARNING:absl:[process=0][thread=MainThread][operation_id=18] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 18 +WARNING:absl:[process=0][thread=MainThread][operation_id=19] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 19 +WARNING:absl:[process=0][thread=MainThread][operation_id=20] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 20 +WARNING:absl:[process=0][thread=MainThread][operation_id=21] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 21 +WARNING:absl:[process=0][thread=MainThread][operation_id=22] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 22 +WARNING:absl:[process=0][thread=MainThread][operation_id=23] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 23 +WARNING:absl:[process=0][thread=MainThread][operation_id=24] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 24 +WARNING:absl:[process=0][thread=MainThread][operation_id=25] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 25 +WARNING:absl:[process=0][thread=MainThread][operation_id=26] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 26 +WARNING:absl:[process=0][thread=MainThread][operation_id=27] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 27 +WARNING:absl:[process=0][thread=MainThread][operation_id=28] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 28 +WARNING:absl:[process=0][thread=MainThread][operation_id=29] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 29 +WARNING:absl:[process=0][thread=MainThread][operation_id=30] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 30 +WARNING:absl:[process=0][thread=MainThread][operation_id=31] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 31 +WARNING:absl:[process=0][thread=MainThread][operation_id=32] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 32 +WARNING:absl:[process=0][thread=MainThread][operation_id=33] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 33 +WARNING:absl:[process=0][thread=MainThread][operation_id=34] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 34 +WARNING:absl:[process=0][thread=MainThread][operation_id=35] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 35 +WARNING:absl:[process=0][thread=MainThread][operation_id=36] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 36 +WARNING:absl:[process=0][thread=MainThread][operation_id=37] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 37 +WARNING:absl:[process=0][thread=MainThread][operation_id=38] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 38 +WARNING:absl:[process=0][thread=MainThread][operation_id=39] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 39 +WARNING:absl:[process=0][thread=MainThread][operation_id=40] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 40 +WARNING:absl:[process=0][thread=MainThread][operation_id=41] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 41 +WARNING:absl:[process=0][thread=MainThread][operation_id=42] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 42 +WARNING:absl:[process=0][thread=MainThread][operation_id=43] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 43 +WARNING:absl:[process=0][thread=MainThread][operation_id=44] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 44 +WARNING:absl:[process=0][thread=MainThread][operation_id=45] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 45 +WARNING:absl:[process=0][thread=MainThread][operation_id=46] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 46 +WARNING:absl:[process=0][thread=MainThread][operation_id=47] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 47 +WARNING:absl:[process=0][thread=MainThread][operation_id=48] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 48 +WARNING:absl:[process=0][thread=MainThread][operation_id=49] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 49 +WARNING:absl:[process=0][thread=MainThread][operation_id=50] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 50 +WARNING:absl:[process=0][thread=MainThread][operation_id=51] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 51 +WARNING:absl:[process=0][thread=MainThread][operation_id=52] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 52 +WARNING:absl:[process=0][thread=MainThread][operation_id=53] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 53 +WARNING:absl:[process=0][thread=MainThread][operation_id=54] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 54 +WARNING:absl:[process=0][thread=MainThread][operation_id=55] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 55 +WARNING:absl:[process=0][thread=MainThread][operation_id=56] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 56 +WARNING:absl:[process=0][thread=MainThread][operation_id=57] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 57 +WARNING:absl:[process=0][thread=MainThread][operation_id=58] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 58 +WARNING:absl:[process=0][thread=MainThread][operation_id=59] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 59 +WARNING:absl:[process=0][thread=MainThread][operation_id=60] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 60 +WARNING:absl:[process=0][thread=MainThread][operation_id=61] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 61 +WARNING:absl:[process=0][thread=MainThread][operation_id=62] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 62 +WARNING:absl:[process=0][thread=MainThread][operation_id=63] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 63 +WARNING:absl:[process=0][thread=MainThread][operation_id=64] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 64 +WARNING:absl:[process=0][thread=MainThread][operation_id=65] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 65 +WARNING:absl:[process=0][thread=MainThread][operation_id=66] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 66 +WARNING:absl:[process=0][thread=MainThread][operation_id=67] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 67 +WARNING:absl:[process=0][thread=MainThread][operation_id=68] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 68 +WARNING:absl:[process=0][thread=MainThread][operation_id=69] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 69 +WARNING:absl:[process=0][thread=MainThread][operation_id=70] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 70 +WARNING:absl:[process=0][thread=MainThread][operation_id=71] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 71 +WARNING:absl:[process=0][thread=MainThread][operation_id=72] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 72 +WARNING:absl:[process=0][thread=MainThread][operation_id=73] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 73 +WARNING:absl:[process=0][thread=MainThread][operation_id=74] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 74 +WARNING:absl:[process=0][thread=MainThread][operation_id=75] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 75 +WARNING:absl:[process=0][thread=MainThread][operation_id=76] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 76 +WARNING:absl:[process=0][thread=MainThread][operation_id=77] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 77 +WARNING:absl:[process=0][thread=MainThread][operation_id=78] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 78 +WARNING:absl:[process=0][thread=MainThread][operation_id=79] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 79 +WARNING:absl:[process=0][thread=MainThread][operation_id=80] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 80 +WARNING:absl:[process=0][thread=MainThread][operation_id=81] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 81 +WARNING:absl:[process=0][thread=MainThread][operation_id=82] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 82 +WARNING:absl:[process=0][thread=MainThread][operation_id=83] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 83 +WARNING:absl:[process=0][thread=MainThread][operation_id=84] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 84 +WARNING:absl:[process=0][thread=MainThread][operation_id=85] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 85 +WARNING:absl:[process=0][thread=MainThread][operation_id=86] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 86 +WARNING:absl:[process=0][thread=MainThread][operation_id=87] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 87 +WARNING:absl:[process=0][thread=MainThread][operation_id=88] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 88 +WARNING:absl:[process=0][thread=MainThread][operation_id=89] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 89 +WARNING:absl:[process=0][thread=MainThread][operation_id=90] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 90 +WARNING:absl:[process=0][thread=MainThread][operation_id=91] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 91 +WARNING:absl:[process=0][thread=MainThread][operation_id=92] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 92 +WARNING:absl:[process=0][thread=MainThread][operation_id=93] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 93 +WARNING:absl:[process=0][thread=MainThread][operation_id=94] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 94 +WARNING:absl:[process=0][thread=MainThread][operation_id=95] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 95 +WARNING:absl:[process=0][thread=MainThread][operation_id=96] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 96 +WARNING:absl:[process=0][thread=MainThread][operation_id=97] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 97 +WARNING:absl:[process=0][thread=MainThread][operation_id=98] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 98 +WARNING:absl:[process=0][thread=MainThread][operation_id=99] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 99 +WARNING:absl:[process=0][thread=MainThread][operation_id=100] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. diff --git a/wandb/run-20251225_200829-ad33gu8w/files/requirements.txt b/wandb/run-20251225_200829-ad33gu8w/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251225_200829-ad33gu8w/files/wandb-metadata.json b/wandb/run-20251225_200829-ad33gu8w/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..af58b492018fca9d7ca9504a341ed0e70e42c2b5 --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/files/wandb-metadata.json @@ -0,0 +1,85 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-25T13:08:29.527349Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "512", + "--num-hidden-layers", + "6", + "--warmup-epochs", + "5", + "--num-attention-heads", + "8", + "--intermediate-size", + "2048", + "--position-embeddings", + "learnable", + "--num-labels", + "100", + "--lr", + "1e-4", + "--epochs", + "100", + "--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", + "/mnt/d/NCKH/weights/lmc/cifar100", + "--data-path", + "/mnt/d/NCKH/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67550371840" + } + }, + "memory": { + "total": "16466989056" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "hxpb1vakuyd64uxdubge65x4g39p6ahi" +} \ No newline at end of file diff --git a/wandb/run-20251225_200829-ad33gu8w/files/wandb-summary.json b/wandb/run-20251225_200829-ad33gu8w/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..c65ff7eeb1f37731147dd823dad9815849bfc3cf --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/files/wandb-summary.json @@ -0,0 +1 @@ +{"lr":1.0001617738453206e-05,"_timestamp":1.7666809978930266e+09,"train/acc1":97.65625,"_step":39000,"train/loss":0.11328125,"_runtime":12879.360673645,"val/acc5":84.375,"_wandb":{"runtime":12879},"val/loss":2.375,"val/acc1":53.125,"train/acc5":100} \ No newline at end of file diff --git a/wandb/run-20251225_200829-ad33gu8w/logs/debug-core.log b/wandb/run-20251225_200829-ad33gu8w/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..8886424a6570ff2f953af8749e8daa19353f4fe3 --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-25T20:08:29.761977808+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmp36ptrs7x/port-426309.txt","pid":426309,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T20:08:29.763934092+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":426309} +{"time":"2025-12-25T20:08:29.76392316+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-426309-426494-3383938975/socket","Net":"unix"}} +{"time":"2025-12-25T20:08:29.858317253+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T20:08:29.8750645+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"ad33gu8w","id":"1(@)"} +{"time":"2025-12-25T20:08:30.421540239+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"ad33gu8w","id":"1(@)"} +{"time":"2025-12-25T23:43:21.363718802+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-25T23:43:21.365412329+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-25T23:43:21.365366225+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-25T23:43:21.365833301+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-25T23:43:21.366837774+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-426309-426494-3383938975/socket","Net":"unix"}} +{"time":"2025-12-25T23:43:22.787058731+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-25T23:43:22.787185913+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-25T23:43:22.78723429+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/run-20251225_200829-ad33gu8w/logs/debug-internal.log b/wandb/run-20251225_200829-ad33gu8w/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..e5d4569dc522f60f08fcd8e833863752ae6c8d2b --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/logs/debug-internal.log @@ -0,0 +1,11 @@ +{"time":"2025-12-25T20:08:29.875258063+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T20:08:30.420671503+07:00","level":"INFO","msg":"stream: created new stream","id":"ad33gu8w"} +{"time":"2025-12-25T20:08:30.421042789+07:00","level":"INFO","msg":"handler: started","stream_id":"ad33gu8w"} +{"time":"2025-12-25T20:08:30.421528903+07:00","level":"INFO","msg":"stream: started","id":"ad33gu8w"} +{"time":"2025-12-25T20:08:30.421676536+07:00","level":"INFO","msg":"writer: started","stream_id":"ad33gu8w"} +{"time":"2025-12-25T20:08:30.421725433+07:00","level":"INFO","msg":"sender: started","stream_id":"ad33gu8w"} +{"time":"2025-12-25T23:43:21.36573672+07:00","level":"INFO","msg":"stream: closing","id":"ad33gu8w"} +{"time":"2025-12-25T23:43:22.378707171+07:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"} +{"time":"2025-12-25T23:43:22.783603006+07:00","level":"INFO","msg":"handler: closed","stream_id":"ad33gu8w"} +{"time":"2025-12-25T23:43:22.783862634+07:00","level":"INFO","msg":"sender: closed","stream_id":"ad33gu8w"} +{"time":"2025-12-25T23:43:22.783883956+07:00","level":"INFO","msg":"stream: closed","id":"ad33gu8w"} diff --git a/wandb/run-20251225_200829-ad33gu8w/logs/debug.log b/wandb/run-20251225_200829-ad33gu8w/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..7d002d4739bb88b6da1cd622f66ffd45026ca509 --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/logs/debug.log @@ -0,0 +1,23 @@ +2025-12-25 20:08:29,533 INFO MainThread:426309 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 20:08:29,533 INFO MainThread:426309 [wandb_setup.py:_flush():80] Configure stats pid to 426309 +2025-12-25 20:08:29,533 INFO MainThread:426309 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 20:08:29,533 INFO MainThread:426309 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 20:08:29,533 INFO MainThread:426309 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 20:08:29,533 INFO MainThread:426309 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_200829-ad33gu8w/logs/debug.log +2025-12-25 20:08:29,533 INFO MainThread:426309 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_200829-ad33gu8w/logs/debug-internal.log +2025-12-25 20:08:29,533 INFO MainThread:426309 [wandb_init.py:init():840] calling init triggers +2025-12-25 20:08:29,533 INFO MainThread:426309 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 20:08:29,533 INFO MainThread:426309 [wandb_init.py:init():888] starting backend +2025-12-25 20:08:29,856 INFO MainThread:426309 [wandb_init.py:init():891] sending inform_init request +2025-12-25 20:08:29,869 INFO MainThread:426309 [wandb_init.py:init():899] backend started and connected +2025-12-25 20:08:29,870 INFO MainThread:426309 [wandb_init.py:init():969] updated telemetry +2025-12-25 20:08:29,876 INFO MainThread:426309 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 20:08:30,852 INFO MainThread:426309 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 20:08:31,003 INFO MainThread:426309 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 20:08:31,004 INFO MainThread:426309 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 20:08:31,004 INFO MainThread:426309 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 20:08:31,004 INFO MainThread:426309 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 20:08:31,007 INFO MainThread:426309 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-25 23:43:21,357 INFO wandb-AsyncioManager-main:426309 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-25 23:43:21,357 INFO wandb-AsyncioManager-main:426309 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. diff --git a/wandb/run-20251225_200829-ad33gu8w/run-ad33gu8w.wandb b/wandb/run-20251225_200829-ad33gu8w/run-ad33gu8w.wandb new file mode 100644 index 0000000000000000000000000000000000000000..8a650f236a0ace00c8e6420fd66a790732d63b25 --- /dev/null +++ b/wandb/run-20251225_200829-ad33gu8w/run-ad33gu8w.wandb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:49801095c01a33af6a5c9366573cd76ff31ef850896450156a0f997ad5441477 +size 10186030 diff --git a/wandb/run-20251225_234450-ddqothxs/files/code/src/imagenet/train_model.py b/wandb/run-20251225_234450-ddqothxs/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c48bfefd7467518c68801e1068bea786a32c1e92 --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/files/code/src/imagenet/train_model.py @@ -0,0 +1,264 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251225_234450-ddqothxs/files/config.yaml b/wandb/run-20251225_234450-ddqothxs/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c217423494bf31b41e13f200420752f40460ca36 --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/files/config.yaml @@ -0,0 +1,116 @@ +_wandb: + value: + cli_version: 0.23.0 + code_path: code/src/imagenet/train_model.py + e: + vhrvck05spgtobo8zo58w4ia49zygoc0: + args: + - --input-size + - "32" + - --data-set + - CIFAR + - --patch-size + - "4" + - --hidden-size + - "512" + - --num-hidden-layers + - "6" + - --warmup-epochs + - "5" + - --num-attention-heads + - "8" + - --intermediate-size + - "2048" + - --position-embeddings + - learnable + - --num-labels + - "100" + - --lr + - "1e-4" + - --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 + - /mnt/d/NCKH/weights/lmc/cifar100 + - --data-path + - /mnt/d/NCKH/datasets/cifar100 + codePath: src/imagenet/train_model.py + codePathLocal: src/imagenet/train_model.py + cpu_count: 24 + cpu_count_logical: 24 + cudaVersion: "13.0" + disk: + /: + total: "1081101176832" + used: "67561218048" + email: khanhvinhbui0512@gmail.com + executable: /home/khanhvinh9/miniconda3/envs/lmc/bin/python + git: + commit: 503490314d9f0bef18a30ec7957b3949bd36c751 + remote: https://github.com/MLResearchX/lmc-transformer.git + gpu: NVIDIA GeForce RTX 5060 Laptop GPU + gpu_count: 1 + gpu_nvidia: + - architecture: Blackwell + memoryTotal: "8546942976" + name: NVIDIA GeForce RTX 5060 Laptop GPU + uuid: GPU-6363d584-7737-137f-cc48-e39ef1a9efbd + host: DESKTOP-UCJ5IKF + memory: + total: "16466989056" + os: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 + program: /home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py + python: CPython 3.10.19 + root: /home/khanhvinh9/lmc-transformer + startedAt: "2025-12-25T16:44:50.877593Z" + writerId: vhrvck05spgtobo8zo58w4ia49zygoc0 + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "2": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "3": + - 13 + - 61 + "4": 3.10.19 + "5": 0.23.0 + "6": 4.56.1 + "8": + - 2 + "12": 0.23.0 + "13": linux-x86_64 diff --git a/wandb/run-20251225_234450-ddqothxs/files/diff.patch b/wandb/run-20251225_234450-ddqothxs/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..64375f8ecec60cacb0096f58826df1afb889a5ae --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/files/diff.patch @@ -0,0 +1,214 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..ffbb4f8 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3f8f560 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +@@ -78,33 +79,48 @@ def build_dataset(is_train, args): + + def build_transform(is_train, args): + resize_im = args.input_size > 32 +- if is_train: +- # this should always dispatch to transforms_imagenet_train +- transform = create_transform( +- input_size=args.input_size, +- is_training=True, +- color_jitter=args.color_jitter, +- auto_augment=args.aa, +- interpolation=args.train_interpolation, +- re_prob=args.reprob, +- re_mode=args.remode, +- re_count=args.recount, +- ) +- if not resize_im: +- # replace RandomResizedCropAndInterpolation with +- # RandomCrop +- transform.transforms[0] = transforms.RandomCrop( +- args.input_size, padding=4) +- return transform +- t = [] +- if resize_im: +- size = int((256 / 224) * args.input_size) +- t.append( +- transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images +- ) +- t.append(transforms.CenterCrop(args.input_size)) +- t.append(transforms.ToTensor()) +- t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) +- return transforms.Compose(t) ++ if args.data_set == 'CIFAR': ++ if is_train: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.RandomHorizontalFlip(), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ if is_train: ++ # this should always dispatch to transforms_imagenet_train ++ transform = create_transform( ++ input_size=args.input_size, ++ is_training=True, ++ color_jitter=args.color_jitter, ++ auto_augment=args.aa, ++ interpolation=args.train_interpolation, ++ re_prob=args.reprob, ++ re_mode=args.remode, ++ re_count=args.recount, ++ ) ++ if not resize_im: ++ # replace RandomResizedCropAndInterpolation with ++ # RandomCrop ++ transform.transforms[0] = transforms.RandomCrop( ++ args.input_size, padding=4) ++ return transform ++ t = [] ++ if resize_im: ++ size = int((256 / 224) * args.input_size) ++ t.append( ++ transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images ++ ) ++ t.append(transforms.CenterCrop(args.input_size)) ++ t.append(transforms.ToTensor()) ++ t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) ++ return transforms.Compose(t) + + +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_234450-ddqothxs/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251225_234450-ddqothxs/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..64375f8ecec60cacb0096f58826df1afb889a5ae --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,214 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..ffbb4f8 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3f8f560 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +@@ -78,33 +79,48 @@ def build_dataset(is_train, args): + + def build_transform(is_train, args): + resize_im = args.input_size > 32 +- if is_train: +- # this should always dispatch to transforms_imagenet_train +- transform = create_transform( +- input_size=args.input_size, +- is_training=True, +- color_jitter=args.color_jitter, +- auto_augment=args.aa, +- interpolation=args.train_interpolation, +- re_prob=args.reprob, +- re_mode=args.remode, +- re_count=args.recount, +- ) +- if not resize_im: +- # replace RandomResizedCropAndInterpolation with +- # RandomCrop +- transform.transforms[0] = transforms.RandomCrop( +- args.input_size, padding=4) +- return transform +- t = [] +- if resize_im: +- size = int((256 / 224) * args.input_size) +- t.append( +- transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images +- ) +- t.append(transforms.CenterCrop(args.input_size)) +- t.append(transforms.ToTensor()) +- t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) +- return transforms.Compose(t) ++ if args.data_set == 'CIFAR': ++ if is_train: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.RandomHorizontalFlip(), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ if is_train: ++ # this should always dispatch to transforms_imagenet_train ++ transform = create_transform( ++ input_size=args.input_size, ++ is_training=True, ++ color_jitter=args.color_jitter, ++ auto_augment=args.aa, ++ interpolation=args.train_interpolation, ++ re_prob=args.reprob, ++ re_mode=args.remode, ++ re_count=args.recount, ++ ) ++ if not resize_im: ++ # replace RandomResizedCropAndInterpolation with ++ # RandomCrop ++ transform.transforms[0] = transforms.RandomCrop( ++ args.input_size, padding=4) ++ return transform ++ t = [] ++ if resize_im: ++ size = int((256 / 224) * args.input_size) ++ t.append( ++ transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images ++ ) ++ t.append(transforms.CenterCrop(args.input_size)) ++ t.append(transforms.ToTensor()) ++ t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) ++ return transforms.Compose(t) + + +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251225_234450-ddqothxs/files/output.log b/wandb/run-20251225_234450-ddqothxs/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..664b9f019d779513893a0cb62a420137d5f19f3a --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/files/output.log @@ -0,0 +1,104 @@ +TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We recommend migrating to PyTorch classes or pinning your version of Transformers. +Starting training... +JAX devices: [CudaDevice(id=0)] +Using 1 devices +Epoch 0 +WARNING:absl:[process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 1 +WARNING:absl:[process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 2 +WARNING:absl:[process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 3 +WARNING:absl:[process=0][thread=MainThread][operation_id=4] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 4 +WARNING:absl:[process=0][thread=MainThread][operation_id=5] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 5 +WARNING:absl:[process=0][thread=MainThread][operation_id=6] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 6 +WARNING:absl:[process=0][thread=MainThread][operation_id=7] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 7 +WARNING:absl:[process=0][thread=MainThread][operation_id=8] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 8 +WARNING:absl:[process=0][thread=MainThread][operation_id=9] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 9 +WARNING:absl:[process=0][thread=MainThread][operation_id=10] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 10 +WARNING:absl:[process=0][thread=MainThread][operation_id=11] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 11 +WARNING:absl:[process=0][thread=MainThread][operation_id=12] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 12 +WARNING:absl:[process=0][thread=MainThread][operation_id=13] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 13 +WARNING:absl:[process=0][thread=MainThread][operation_id=14] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 14 +WARNING:absl:[process=0][thread=MainThread][operation_id=15] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 15 +WARNING:absl:[process=0][thread=MainThread][operation_id=16] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 16 +WARNING:absl:[process=0][thread=MainThread][operation_id=17] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 17 +WARNING:absl:[process=0][thread=MainThread][operation_id=18] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 18 +WARNING:absl:[process=0][thread=MainThread][operation_id=19] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 19 +WARNING:absl:[process=0][thread=MainThread][operation_id=20] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 20 +WARNING:absl:[process=0][thread=MainThread][operation_id=21] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 21 +WARNING:absl:[process=0][thread=MainThread][operation_id=22] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 22 +WARNING:absl:[process=0][thread=MainThread][operation_id=23] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 23 +WARNING:absl:[process=0][thread=MainThread][operation_id=24] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 24 +WARNING:absl:[process=0][thread=MainThread][operation_id=25] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 25 +WARNING:absl:[process=0][thread=MainThread][operation_id=26] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 26 +WARNING:absl:[process=0][thread=MainThread][operation_id=27] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 27 +WARNING:absl:[process=0][thread=MainThread][operation_id=28] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 28 +WARNING:absl:[process=0][thread=MainThread][operation_id=29] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 29 +WARNING:absl:[process=0][thread=MainThread][operation_id=30] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 30 +WARNING:absl:[process=0][thread=MainThread][operation_id=31] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 31 +WARNING:absl:[process=0][thread=MainThread][operation_id=32] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 32 +WARNING:absl:[process=0][thread=MainThread][operation_id=33] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 33 +WARNING:absl:[process=0][thread=MainThread][operation_id=34] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 34 +WARNING:absl:[process=0][thread=MainThread][operation_id=35] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 35 +WARNING:absl:[process=0][thread=MainThread][operation_id=36] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 36 +WARNING:absl:[process=0][thread=MainThread][operation_id=37] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 37 +WARNING:absl:[process=0][thread=MainThread][operation_id=38] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 38 +WARNING:absl:[process=0][thread=MainThread][operation_id=39] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 39 +WARNING:absl:[process=0][thread=MainThread][operation_id=40] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 40 +WARNING:absl:[process=0][thread=MainThread][operation_id=41] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 41 +WARNING:absl:[process=0][thread=MainThread][operation_id=42] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 42 +WARNING:absl:[process=0][thread=MainThread][operation_id=43] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 43 +WARNING:absl:[process=0][thread=MainThread][operation_id=44] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 44 +WARNING:absl:[process=0][thread=MainThread][operation_id=45] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 45 +WARNING:absl:[process=0][thread=MainThread][operation_id=46] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 46 +WARNING:absl:[process=0][thread=MainThread][operation_id=47] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 47 +WARNING:absl:[process=0][thread=MainThread][operation_id=48] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 48 +WARNING:absl:[process=0][thread=MainThread][operation_id=49] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 49 +WARNING:absl:[process=0][thread=MainThread][operation_id=50] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. diff --git a/wandb/run-20251225_234450-ddqothxs/files/requirements.txt b/wandb/run-20251225_234450-ddqothxs/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251225_234450-ddqothxs/files/wandb-metadata.json b/wandb/run-20251225_234450-ddqothxs/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..c6179747814fc3c0a69cb0e41dfcd251056d2caa --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/files/wandb-metadata.json @@ -0,0 +1,85 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-25T16:44:50.877593Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "512", + "--num-hidden-layers", + "6", + "--warmup-epochs", + "5", + "--num-attention-heads", + "8", + "--intermediate-size", + "2048", + "--position-embeddings", + "learnable", + "--num-labels", + "100", + "--lr", + "1e-4", + "--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", + "/mnt/d/NCKH/weights/lmc/cifar100", + "--data-path", + "/mnt/d/NCKH/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67561218048" + } + }, + "memory": { + "total": "16466989056" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "vhrvck05spgtobo8zo58w4ia49zygoc0" +} \ No newline at end of file diff --git a/wandb/run-20251225_234450-ddqothxs/files/wandb-summary.json b/wandb/run-20251225_234450-ddqothxs/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..5d457714238dfb22286741dd4e691ff7925b6f08 --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/files/wandb-summary.json @@ -0,0 +1 @@ +{"train/loss":0.9765625,"val/loss":1.765625,"lr":1.0007210221374407e-05,"_step":19500,"val/acc5":82.03125,"_runtime":6737.498793088,"train/acc5":95.3125,"_wandb":{"runtime":6737},"_timestamp":1.7666878268165345e+09,"val/acc1":59.375,"train/acc1":70.3125} \ No newline at end of file diff --git a/wandb/run-20251225_234450-ddqothxs/logs/debug-core.log b/wandb/run-20251225_234450-ddqothxs/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..f7e1eec3fc63ec36410052397c5aaa6030ad2e89 --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-25T23:44:50.969406686+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmpl053p3mo/port-767046.txt","pid":767046,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-25T23:44:50.970448212+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":767046} +{"time":"2025-12-25T23:44:50.970408717+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-767046-767165-4282731576/socket","Net":"unix"}} +{"time":"2025-12-25T23:44:51.141214786+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-25T23:44:51.148546504+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"ddqothxs","id":"1(@)"} +{"time":"2025-12-25T23:44:51.583115185+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"ddqothxs","id":"1(@)"} +{"time":"2025-12-26T01:37:09.426824557+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-26T01:37:09.427266489+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-26T01:37:09.427212745+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-26T01:37:09.427445349+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-26T01:37:09.427664832+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-767046-767165-4282731576/socket","Net":"unix"}} +{"time":"2025-12-26T01:37:11.417716916+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-26T01:37:11.417822466+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-26T01:37:11.418011234+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/run-20251225_234450-ddqothxs/logs/debug-internal.log b/wandb/run-20251225_234450-ddqothxs/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..e0938b8fedd4720e97bf3040983f54db590bdd6a --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/logs/debug-internal.log @@ -0,0 +1,11 @@ +{"time":"2025-12-25T23:44:51.148670994+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-25T23:44:51.58275059+07:00","level":"INFO","msg":"stream: created new stream","id":"ddqothxs"} +{"time":"2025-12-25T23:44:51.582818752+07:00","level":"INFO","msg":"handler: started","stream_id":"ddqothxs"} +{"time":"2025-12-25T23:44:51.583107675+07:00","level":"INFO","msg":"stream: started","id":"ddqothxs"} +{"time":"2025-12-25T23:44:51.583136474+07:00","level":"INFO","msg":"writer: started","stream_id":"ddqothxs"} +{"time":"2025-12-25T23:44:51.583151546+07:00","level":"INFO","msg":"sender: started","stream_id":"ddqothxs"} +{"time":"2025-12-26T01:37:09.427299622+07:00","level":"INFO","msg":"stream: closing","id":"ddqothxs"} +{"time":"2025-12-26T01:37:10.859592581+07:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"} +{"time":"2025-12-26T01:37:11.413068085+07:00","level":"INFO","msg":"handler: closed","stream_id":"ddqothxs"} +{"time":"2025-12-26T01:37:11.413242832+07:00","level":"INFO","msg":"sender: closed","stream_id":"ddqothxs"} +{"time":"2025-12-26T01:37:11.413254953+07:00","level":"INFO","msg":"stream: closed","id":"ddqothxs"} diff --git a/wandb/run-20251225_234450-ddqothxs/logs/debug.log b/wandb/run-20251225_234450-ddqothxs/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..2c03372adb54c580eb942ea85ca5de3b4604efd0 --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/logs/debug.log @@ -0,0 +1,23 @@ +2025-12-25 23:44:50,879 INFO MainThread:767046 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-25 23:44:50,879 INFO MainThread:767046 [wandb_setup.py:_flush():80] Configure stats pid to 767046 +2025-12-25 23:44:50,879 INFO MainThread:767046 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-25 23:44:50,879 INFO MainThread:767046 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-25 23:44:50,879 INFO MainThread:767046 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-25 23:44:50,879 INFO MainThread:767046 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_234450-ddqothxs/logs/debug.log +2025-12-25 23:44:50,879 INFO MainThread:767046 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251225_234450-ddqothxs/logs/debug-internal.log +2025-12-25 23:44:50,879 INFO MainThread:767046 [wandb_init.py:init():840] calling init triggers +2025-12-25 23:44:50,879 INFO MainThread:767046 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-25 23:44:50,879 INFO MainThread:767046 [wandb_init.py:init():888] starting backend +2025-12-25 23:44:51,140 INFO MainThread:767046 [wandb_init.py:init():891] sending inform_init request +2025-12-25 23:44:51,144 INFO MainThread:767046 [wandb_init.py:init():899] backend started and connected +2025-12-25 23:44:51,145 INFO MainThread:767046 [wandb_init.py:init():969] updated telemetry +2025-12-25 23:44:51,149 INFO MainThread:767046 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-25 23:44:51,928 INFO MainThread:767046 [wandb_init.py:init():1040] starting run threads in backend +2025-12-25 23:44:51,983 INFO MainThread:767046 [wandb_run.py:_console_start():2504] atexit reg +2025-12-25 23:44:51,983 INFO MainThread:767046 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-25 23:44:51,983 INFO MainThread:767046 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-25 23:44:51,983 INFO MainThread:767046 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-25 23:44:51,984 INFO MainThread:767046 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-26 01:37:09,425 INFO wandb-AsyncioManager-main:767046 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-26 01:37:09,425 INFO wandb-AsyncioManager-main:767046 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. diff --git a/wandb/run-20251225_234450-ddqothxs/run-ddqothxs.wandb b/wandb/run-20251225_234450-ddqothxs/run-ddqothxs.wandb new file mode 100644 index 0000000000000000000000000000000000000000..6a44af388aa2679ace1febc4d0bdb34e0a3ba111 --- /dev/null +++ b/wandb/run-20251225_234450-ddqothxs/run-ddqothxs.wandb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:483e589cebad74a3faa71bb087416e6c94368c07c34b0db558f8bda1e891dc6b +size 5087409 diff --git a/wandb/run-20251226_013720-ly1ggf35/files/code/src/imagenet/train_model.py b/wandb/run-20251226_013720-ly1ggf35/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c48bfefd7467518c68801e1068bea786a32c1e92 --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/files/code/src/imagenet/train_model.py @@ -0,0 +1,264 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251226_013720-ly1ggf35/files/config.yaml b/wandb/run-20251226_013720-ly1ggf35/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..21e1f50ff689e30ebb4ad6081c11da540f7ec724 --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/files/config.yaml @@ -0,0 +1,116 @@ +_wandb: + value: + cli_version: 0.23.0 + code_path: code/src/imagenet/train_model.py + e: + 2jein137km7ssb4nxvf7b8c6ya0zxqf3: + args: + - --input-size + - "32" + - --data-set + - CIFAR + - --patch-size + - "4" + - --hidden-size + - "384" + - --num-hidden-layers + - "6" + - --warmup-epochs + - "5" + - --num-attention-heads + - "6" + - --intermediate-size + - "1536" + - --position-embeddings + - learnable + - --num-labels + - "100" + - --lr + - "1e-4" + - --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 + - /mnt/d/NCKH/weights/lmc/cifar100 + - --data-path + - /mnt/d/NCKH/datasets/cifar100 + codePath: src/imagenet/train_model.py + codePathLocal: src/imagenet/train_model.py + cpu_count: 24 + cpu_count_logical: 24 + cudaVersion: "13.0" + disk: + /: + total: "1081101176832" + used: "67566546944" + email: khanhvinhbui0512@gmail.com + executable: /home/khanhvinh9/miniconda3/envs/lmc/bin/python + git: + commit: 503490314d9f0bef18a30ec7957b3949bd36c751 + remote: https://github.com/MLResearchX/lmc-transformer.git + gpu: NVIDIA GeForce RTX 5060 Laptop GPU + gpu_count: 1 + gpu_nvidia: + - architecture: Blackwell + memoryTotal: "8546942976" + name: NVIDIA GeForce RTX 5060 Laptop GPU + uuid: GPU-6363d584-7737-137f-cc48-e39ef1a9efbd + host: DESKTOP-UCJ5IKF + memory: + total: "16466989056" + os: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 + program: /home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py + python: CPython 3.10.19 + root: /home/khanhvinh9/lmc-transformer + startedAt: "2025-12-25T18:37:20.380252Z" + writerId: 2jein137km7ssb4nxvf7b8c6ya0zxqf3 + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "2": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "3": + - 13 + - 61 + "4": 3.10.19 + "5": 0.23.0 + "6": 4.56.1 + "8": + - 2 + "12": 0.23.0 + "13": linux-x86_64 diff --git a/wandb/run-20251226_013720-ly1ggf35/files/diff.patch b/wandb/run-20251226_013720-ly1ggf35/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..64375f8ecec60cacb0096f58826df1afb889a5ae --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/files/diff.patch @@ -0,0 +1,214 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..ffbb4f8 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3f8f560 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +@@ -78,33 +79,48 @@ def build_dataset(is_train, args): + + def build_transform(is_train, args): + resize_im = args.input_size > 32 +- if is_train: +- # this should always dispatch to transforms_imagenet_train +- transform = create_transform( +- input_size=args.input_size, +- is_training=True, +- color_jitter=args.color_jitter, +- auto_augment=args.aa, +- interpolation=args.train_interpolation, +- re_prob=args.reprob, +- re_mode=args.remode, +- re_count=args.recount, +- ) +- if not resize_im: +- # replace RandomResizedCropAndInterpolation with +- # RandomCrop +- transform.transforms[0] = transforms.RandomCrop( +- args.input_size, padding=4) +- return transform +- t = [] +- if resize_im: +- size = int((256 / 224) * args.input_size) +- t.append( +- transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images +- ) +- t.append(transforms.CenterCrop(args.input_size)) +- t.append(transforms.ToTensor()) +- t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) +- return transforms.Compose(t) ++ if args.data_set == 'CIFAR': ++ if is_train: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.RandomHorizontalFlip(), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ if is_train: ++ # this should always dispatch to transforms_imagenet_train ++ transform = create_transform( ++ input_size=args.input_size, ++ is_training=True, ++ color_jitter=args.color_jitter, ++ auto_augment=args.aa, ++ interpolation=args.train_interpolation, ++ re_prob=args.reprob, ++ re_mode=args.remode, ++ re_count=args.recount, ++ ) ++ if not resize_im: ++ # replace RandomResizedCropAndInterpolation with ++ # RandomCrop ++ transform.transforms[0] = transforms.RandomCrop( ++ args.input_size, padding=4) ++ return transform ++ t = [] ++ if resize_im: ++ size = int((256 / 224) * args.input_size) ++ t.append( ++ transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images ++ ) ++ t.append(transforms.CenterCrop(args.input_size)) ++ t.append(transforms.ToTensor()) ++ t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) ++ return transforms.Compose(t) + + +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251226_013720-ly1ggf35/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251226_013720-ly1ggf35/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..64375f8ecec60cacb0096f58826df1afb889a5ae --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,214 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..ffbb4f8 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3f8f560 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +@@ -78,33 +79,48 @@ def build_dataset(is_train, args): + + def build_transform(is_train, args): + resize_im = args.input_size > 32 +- if is_train: +- # this should always dispatch to transforms_imagenet_train +- transform = create_transform( +- input_size=args.input_size, +- is_training=True, +- color_jitter=args.color_jitter, +- auto_augment=args.aa, +- interpolation=args.train_interpolation, +- re_prob=args.reprob, +- re_mode=args.remode, +- re_count=args.recount, +- ) +- if not resize_im: +- # replace RandomResizedCropAndInterpolation with +- # RandomCrop +- transform.transforms[0] = transforms.RandomCrop( +- args.input_size, padding=4) +- return transform +- t = [] +- if resize_im: +- size = int((256 / 224) * args.input_size) +- t.append( +- transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images +- ) +- t.append(transforms.CenterCrop(args.input_size)) +- t.append(transforms.ToTensor()) +- t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) +- return transforms.Compose(t) ++ if args.data_set == 'CIFAR': ++ if is_train: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.RandomHorizontalFlip(), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ if is_train: ++ # this should always dispatch to transforms_imagenet_train ++ transform = create_transform( ++ input_size=args.input_size, ++ is_training=True, ++ color_jitter=args.color_jitter, ++ auto_augment=args.aa, ++ interpolation=args.train_interpolation, ++ re_prob=args.reprob, ++ re_mode=args.remode, ++ re_count=args.recount, ++ ) ++ if not resize_im: ++ # replace RandomResizedCropAndInterpolation with ++ # RandomCrop ++ transform.transforms[0] = transforms.RandomCrop( ++ args.input_size, padding=4) ++ return transform ++ t = [] ++ if resize_im: ++ size = int((256 / 224) * args.input_size) ++ t.append( ++ transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images ++ ) ++ t.append(transforms.CenterCrop(args.input_size)) ++ t.append(transforms.ToTensor()) ++ t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) ++ return transforms.Compose(t) + + +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251226_013720-ly1ggf35/files/output.log b/wandb/run-20251226_013720-ly1ggf35/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..3dc37dd3d3ae1da352c3cf340e4bc199dd63f96f --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/files/output.log @@ -0,0 +1,123 @@ +TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We recommend migrating to PyTorch classes or pinning your version of Transformers. +Starting training... +JAX devices: [CudaDevice(id=0)] +Using 1 devices +Epoch 0 +WARNING:absl:[process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 1 +WARNING:absl:[process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 2 +WARNING:absl:[process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 3 +WARNING:absl:[process=0][thread=MainThread][operation_id=4] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 4 +WARNING:absl:[process=0][thread=MainThread][operation_id=5] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 5 +WARNING:absl:[process=0][thread=MainThread][operation_id=6] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 6 +WARNING:absl:[process=0][thread=MainThread][operation_id=7] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 7 +WARNING:absl:[process=0][thread=MainThread][operation_id=8] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 8 +WARNING:absl:[process=0][thread=MainThread][operation_id=9] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 9 +WARNING:absl:[process=0][thread=MainThread][operation_id=10] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 10 +WARNING:absl:[process=0][thread=MainThread][operation_id=11] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 11 +WARNING:absl:[process=0][thread=MainThread][operation_id=12] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 12 +WARNING:absl:[process=0][thread=MainThread][operation_id=13] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 13 +WARNING:absl:[process=0][thread=MainThread][operation_id=14] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 14 +WARNING:absl:[process=0][thread=MainThread][operation_id=15] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 15 +WARNING:absl:[process=0][thread=MainThread][operation_id=16] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 16 +WARNING:absl:[process=0][thread=MainThread][operation_id=17] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 17 +WARNING:absl:[process=0][thread=MainThread][operation_id=18] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 18 +WARNING:absl:[process=0][thread=MainThread][operation_id=19] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 19 +WARNING:absl:[process=0][thread=MainThread][operation_id=20] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 20 +WARNING:absl:[process=0][thread=MainThread][operation_id=21] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 21 +WARNING:absl:[process=0][thread=MainThread][operation_id=22] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 22 +WARNING:absl:[process=0][thread=MainThread][operation_id=23] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 23 +WARNING:absl:[process=0][thread=MainThread][operation_id=24] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 24 +WARNING:absl:[process=0][thread=MainThread][operation_id=25] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 25 +WARNING:absl:[process=0][thread=MainThread][operation_id=26] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 26 +WARNING:absl:[process=0][thread=MainThread][operation_id=27] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 27 +WARNING:absl:[process=0][thread=MainThread][operation_id=28] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 28 +WARNING:absl:[process=0][thread=MainThread][operation_id=29] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 29 +WARNING:absl:[process=0][thread=MainThread][operation_id=30] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 30 +WARNING:absl:[process=0][thread=MainThread][operation_id=31] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 31 +WARNING:absl:[process=0][thread=MainThread][operation_id=32] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 32 +WARNING:absl:[process=0][thread=MainThread][operation_id=33] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 33 +WARNING:absl:[process=0][thread=MainThread][operation_id=34] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 34 +WARNING:absl:[process=0][thread=MainThread][operation_id=35] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 35 +WARNING:absl:[process=0][thread=MainThread][operation_id=36] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 36 +WARNING:absl:[process=0][thread=MainThread][operation_id=37] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 37 +WARNING:absl:[process=0][thread=MainThread][operation_id=38] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 38 +WARNING:absl:[process=0][thread=MainThread][operation_id=39] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 39 +WARNING:absl:[process=0][thread=MainThread][operation_id=40] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 40 +WARNING:absl:[process=0][thread=MainThread][operation_id=41] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 41 +WARNING:absl:[process=0][thread=MainThread][operation_id=42] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 42 +WARNING:absl:[process=0][thread=MainThread][operation_id=43] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 43 +WARNING:absl:[process=0][thread=MainThread][operation_id=44] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 44 +WARNING:absl:[process=0][thread=MainThread][operation_id=45] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 45 +WARNING:absl:[process=0][thread=MainThread][operation_id=46] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 46 +WARNING:absl:[process=0][thread=MainThread][operation_id=47] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 47 +WARNING:absl:[process=0][thread=MainThread][operation_id=48] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 48 +WARNING:absl:[process=0][thread=MainThread][operation_id=49] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture. +Epoch 49 +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 207, in main + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/flax/training/checkpoints.py", line 693, in save_checkpoint + orbax_checkpointer.save( + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 254, in save + raise ValueError(f'Destination {directory} already exists.') +ValueError: Destination /mnt/d/NCKH/weights/lmc/cifar100/lr0.0001-learnable-epochs50-batch128-seed0/last_19500 already exists. +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 207, in main + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/flax/training/checkpoints.py", line 693, in save_checkpoint + orbax_checkpointer.save( + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 254, in save + raise ValueError(f'Destination {directory} already exists.') +ValueError: Destination /mnt/d/NCKH/weights/lmc/cifar100/lr0.0001-learnable-epochs50-batch128-seed0/last_19500 already exists. diff --git a/wandb/run-20251226_013720-ly1ggf35/files/requirements.txt b/wandb/run-20251226_013720-ly1ggf35/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251226_013720-ly1ggf35/files/wandb-metadata.json b/wandb/run-20251226_013720-ly1ggf35/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..bda08c9b57a4a37804829563d515930879120c0d --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/files/wandb-metadata.json @@ -0,0 +1,85 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-25T18:37:20.380252Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "384", + "--num-hidden-layers", + "6", + "--warmup-epochs", + "5", + "--num-attention-heads", + "6", + "--intermediate-size", + "1536", + "--position-embeddings", + "learnable", + "--num-labels", + "100", + "--lr", + "1e-4", + "--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", + "/mnt/d/NCKH/weights/lmc/cifar100", + "--data-path", + "/mnt/d/NCKH/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67566546944" + } + }, + "memory": { + "total": "16466989056" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "2jein137km7ssb4nxvf7b8c6ya0zxqf3" +} \ No newline at end of file diff --git a/wandb/run-20251226_013720-ly1ggf35/files/wandb-summary.json b/wandb/run-20251226_013720-ly1ggf35/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..c40159c1895c8e1ef21ada76a939fd95ac63314d --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/files/wandb-summary.json @@ -0,0 +1 @@ +{"val/acc1":53.90625,"val/acc5":84.375,"train/loss":1.3671875,"train/acc5":85.9375,"_step":19500,"_wandb":{"runtime":5930},"_timestamp":1.7666937727101853e+09,"val/loss":1.6640625,"_runtime":5930.835574035,"train/acc1":60.9375,"lr":1.0007210221374407e-05} \ No newline at end of file diff --git a/wandb/run-20251226_013720-ly1ggf35/logs/debug-core.log b/wandb/run-20251226_013720-ly1ggf35/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..1bad3ac5e56c7e35cb762ad6e636edfdff4f3b61 --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/logs/debug-core.log @@ -0,0 +1,14 @@ +{"time":"2025-12-26T01:37:20.528382168+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmp95j9v88u/port-948822.txt","pid":948822,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-26T01:37:20.529667577+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":948822} +{"time":"2025-12-26T01:37:20.529615715+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-948822-948947-3152241348/socket","Net":"unix"}} +{"time":"2025-12-26T01:37:20.659975926+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-26T01:37:20.670874841+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"ly1ggf35","id":"1(@)"} +{"time":"2025-12-26T01:37:21.117516719+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"ly1ggf35","id":"1(@)"} +{"time":"2025-12-26T03:16:12.750162339+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-26T03:16:12.750298623+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-26T03:16:12.750479058+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-26T03:16:12.75032311+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-26T03:16:12.750613018+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-948822-948947-3152241348/socket","Net":"unix"}} +{"time":"2025-12-26T03:16:14.651931487+07:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"} +{"time":"2025-12-26T03:16:14.651983323+07:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"} +{"time":"2025-12-26T03:16:14.651994994+07:00","level":"INFO","msg":"server is closed"} diff --git a/wandb/run-20251226_013720-ly1ggf35/logs/debug-internal.log b/wandb/run-20251226_013720-ly1ggf35/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..88c01c93be30a0fb071e98fca22fc525479cc722 --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/logs/debug-internal.log @@ -0,0 +1,11 @@ +{"time":"2025-12-26T01:37:20.671083767+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-26T01:37:21.117049373+07:00","level":"INFO","msg":"stream: created new stream","id":"ly1ggf35"} +{"time":"2025-12-26T01:37:21.117243214+07:00","level":"INFO","msg":"handler: started","stream_id":"ly1ggf35"} +{"time":"2025-12-26T01:37:21.117466871+07:00","level":"INFO","msg":"writer: started","stream_id":"ly1ggf35"} +{"time":"2025-12-26T01:37:21.117487202+07:00","level":"INFO","msg":"stream: started","id":"ly1ggf35"} +{"time":"2025-12-26T01:37:21.117491908+07:00","level":"INFO","msg":"sender: started","stream_id":"ly1ggf35"} +{"time":"2025-12-26T03:16:12.750256028+07:00","level":"INFO","msg":"stream: closing","id":"ly1ggf35"} +{"time":"2025-12-26T03:16:14.183871848+07:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"} +{"time":"2025-12-26T03:16:14.648748562+07:00","level":"INFO","msg":"handler: closed","stream_id":"ly1ggf35"} +{"time":"2025-12-26T03:16:14.648951654+07:00","level":"INFO","msg":"sender: closed","stream_id":"ly1ggf35"} +{"time":"2025-12-26T03:16:14.648966421+07:00","level":"INFO","msg":"stream: closed","id":"ly1ggf35"} diff --git a/wandb/run-20251226_013720-ly1ggf35/logs/debug.log b/wandb/run-20251226_013720-ly1ggf35/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..2c39b86d99353b6c6442176d5648c3383f930f48 --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/logs/debug.log @@ -0,0 +1,23 @@ +2025-12-26 01:37:20,382 INFO MainThread:948822 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-26 01:37:20,383 INFO MainThread:948822 [wandb_setup.py:_flush():80] Configure stats pid to 948822 +2025-12-26 01:37:20,383 INFO MainThread:948822 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-26 01:37:20,383 INFO MainThread:948822 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-26 01:37:20,383 INFO MainThread:948822 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-26 01:37:20,383 INFO MainThread:948822 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251226_013720-ly1ggf35/logs/debug.log +2025-12-26 01:37:20,383 INFO MainThread:948822 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251226_013720-ly1ggf35/logs/debug-internal.log +2025-12-26 01:37:20,383 INFO MainThread:948822 [wandb_init.py:init():840] calling init triggers +2025-12-26 01:37:20,383 INFO MainThread:948822 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-26 01:37:20,383 INFO MainThread:948822 [wandb_init.py:init():888] starting backend +2025-12-26 01:37:20,660 INFO MainThread:948822 [wandb_init.py:init():891] sending inform_init request +2025-12-26 01:37:20,666 INFO MainThread:948822 [wandb_init.py:init():899] backend started and connected +2025-12-26 01:37:20,667 INFO MainThread:948822 [wandb_init.py:init():969] updated telemetry +2025-12-26 01:37:20,675 INFO MainThread:948822 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-26 01:37:21,507 INFO MainThread:948822 [wandb_init.py:init():1040] starting run threads in backend +2025-12-26 01:37:21,632 INFO MainThread:948822 [wandb_run.py:_console_start():2504] atexit reg +2025-12-26 01:37:21,633 INFO MainThread:948822 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-26 01:37:21,633 INFO MainThread:948822 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-26 01:37:21,633 INFO MainThread:948822 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-26 01:37:21,635 INFO MainThread:948822 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-26 03:16:12,749 INFO wandb-AsyncioManager-main:948822 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-26 03:16:12,749 INFO wandb-AsyncioManager-main:948822 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. diff --git a/wandb/run-20251226_013720-ly1ggf35/run-ly1ggf35.wandb b/wandb/run-20251226_013720-ly1ggf35/run-ly1ggf35.wandb new file mode 100644 index 0000000000000000000000000000000000000000..f4811275e5171ddb9d5371e52c7f9d67ac6e4595 --- /dev/null +++ b/wandb/run-20251226_013720-ly1ggf35/run-ly1ggf35.wandb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b2e4df4f723b8774f2a606f5ed6c65e09a1b47d6ebe77500e4e34790ef595ae1 +size 4904089 diff --git a/wandb/run-20251226_094108-svj9s262/files/code/src/imagenet/train_model.py b/wandb/run-20251226_094108-svj9s262/files/code/src/imagenet/train_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c48bfefd7467518c68801e1068bea786a32c1e92 --- /dev/null +++ b/wandb/run-20251226_094108-svj9s262/files/code/src/imagenet/train_model.py @@ -0,0 +1,264 @@ +import argparse +import os +import jax +import flax +import copy +import json +import optax +import torch +import wandb +import numpy as np +from tqdm import tqdm +import jax.numpy as jnp +from flax import linen as nn +from flax.jax_utils import replicate, unreplicate +from flax.training import checkpoints, train_state +from flax.core.frozen_dict import freeze, unfreeze +from flax.traverse_util import flatten_dict, unflatten_dict +from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification +from flax.training.common_utils import get_metrics, onehot, shard +from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix +from datasets import build_dataset +import multiprocessing as mp +from jax import debug +from pprint import pprint +from typing import Any, Dict, List +import shutil +mp.set_start_method("spawn", force=True) +os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" +def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() + for key in metrics_list[0] + } +# ---------- Dataset Loader ---------- +def remove_old_dirs_with_prefix(save_path, prefix, keep_step): + for fname in os.listdir(save_path): + if fname.startswith(prefix) and not fname.endswith(str(keep_step)): + full_path = os.path.join(save_path, fname) + if os.path.isdir(full_path): + shutil.rmtree(full_path) +def imagenet_data_loader(args): + dataset_train, args.nb_classes = build_dataset(is_train=True, args=args) + dataset_val, _ = build_dataset(is_train=False, args=args) + sampler_train = torch.utils.data.RandomSampler(dataset_train) + sampler_val = torch.utils.data.SequentialSampler(dataset_val) + data_loader_train = torch.utils.data.DataLoader( + dataset_train, sampler=sampler_train, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True + ) + data_loader_val = torch.utils.data.DataLoader( + dataset_val, sampler=sampler_val, batch_size=args.batch_size, + num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False + ) + return data_loader_train, data_loader_val +# fmt: on +def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]: + images, labels = jnp.array(images),jnp.array(labels) + return {'images': shard(images),'labels': shard(labels)} +def accuracy(logits, labels, topk=(1,)): + maxk = max(topk) + batch_size = labels.shape[0] + topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1] # Top-k predictions + res = [] + for k in topk: + correct = (topk_preds[:, :k] == labels[:, None]) + correct = jnp.any(correct, axis=1) + correct = jnp.sum(correct) + res.append(100.0 * correct / batch_size) + return res # list of [acc@1, acc@5] +def main(args: argparse.Namespace): + # --- Seeds & RNG --- + wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, + name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) + save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + rng = jax.random.PRNGKey(args.seed) + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- + config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size + config.patch_size = args.patch_size + config.image_size = args.input_size + config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts + config.num_routed_experts = args.num_routed_experts + config.topk = args.topk + config.routed_scaling_factor = args.routed_scaling_factor + config.lmc_layer_indices = args.lmc_layer_indices + model = LMCFlaxViTForImageClassification( + config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype), + ) + model.config.save_pretrained(save_path) + num_global_steps = len(train_loader)*args.epochs + num_warmup_steps = len(train_loader)*args.warmup_epochs + lr_schedule =optax.warmup_cosine_decay_schedule( + init_value=args.warmup_lr, + peak_value=args.lr, + warmup_steps=num_warmup_steps, + decay_steps=num_global_steps, + end_value=args.min_lr, + ) + tx = optax.adamw( + learning_rate=lr_schedule, + b1=args.adamw_beta1, + b2=args.adamw_beta2, + eps=args.adamw_eps, + weight_decay=args.weight_decay, + ) + state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx) + if args.restore_checkpoint_path: + state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state) + print(f"train state restored from {args.restore_checkpoint_path}") + print(f"skip trian step to {state.step}") + latest_global_step = state.step + curr_epoch = latest_global_step//len(train_loader) + state = replicate(state) + def train_step(state, batch, rng): + dropout_rng, new_dropout_rng = jax.random.split(rng) + def loss_fn(params): + outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + return loss, logits + (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) + grads = jax.lax.pmean(grads, axis_name="batch") + state = state.apply_gradients(grads=grads) + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return state, metrics, new_dropout_rng + def eval_step(state, batch): + outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean() + acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5)) + metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,} + metrics = jax.lax.pmean(metrics, axis_name="batch") + return metrics + parallel_train_step = jax.pmap(train_step, "batch") + parallel_eval_step = jax.pmap(eval_step, "batch") + rng = jax.random.PRNGKey(args.seed) + global_step = latest_global_step + best_val_acc1 = 85.9375 + print("Starting training...") + print(f"JAX devices: {jax.devices()}") + print(f"Using {jax.local_device_count()} devices") + for epoch in range(curr_epoch,args.epochs): + print(f"Epoch {epoch}") + dropout_rngs = jax.random.split(rng, jax.local_device_count()) + train_metrics_stack = [] + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + for batch_idx, (images, labels) in pbar: + # Prepare and shard batch + batch = prepare_image_batch(images,labels) + # Run train step + state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs) + train_metrics = unreplicate(train_metrics) + train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics) + loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),) + curr_lr = float(lr_schedule(global_step)) + if global_step % args.wandb_logging_frequency == 0: + wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step) + pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",}) + global_step += 1 + eval_results = [] + pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False) + for batch_idx, (images, labels) in pbar: + batch = prepare_image_batch(images, labels) + eval_metric = parallel_eval_step(state, batch) + eval_results.append(eval_metric) + # Compute mean metrics across all eval batches + eval_metrics = get_metrics(eval_results) + eval_metrics = unreplicate(eval_metrics) + eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) + val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"]) + wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step) + # Save best model + if best_val_acc1 <= val_acc1: + best_val_acc1 = val_acc1 + # model.params = unreplicate(state).params + # best_dir = os.path.join(save_path, f"best_{global_step}") + # model.save_pretrained(best_dir) + # print(f"Best model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "best_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1) + # Save last model + # model.params = unreplicate(state).params + # last_dir = os.path.join(save_path, f"last_{global_step}") + # model.save_pretrained(last_dir) + # print(f"Last model saved at step {global_step}") + # remove_old_dirs_with_prefix(save_path, "last_", global_step) + checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet") + # --- Model & Training Config --- + parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.") + parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.") + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') + parser.add_argument("--patch-size", type=int, default=16, help="") + parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) + parser.add_argument("--topk", type=int, default = 0) + parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.') + parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.') + parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.') + parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.') + parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.') + parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.') + parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='') + parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)") + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER') + parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') + parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports') + parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') + parser.add_argument("--adamw-beta1", type=float, default=0.9) + parser.add_argument("--adamw-beta2", type=float, default=0.999) + parser.add_argument("--adamw-eps", type=float, default=1e-8) + parser.add_argument("--patience", type=int, default=10, help="Early stopping patience") + parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path") + # --- Data Config --- + parser.add_argument("--data-path", type=str, required=True) + parser.add_argument("--save-dir", type=str, required=True) + parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19']) + parser.add_argument("--input-size", type=int, default=224) + parser.add_argument('--num_workers', type=int, default=10) + parser.add_argument('--pin-mem', action='store_true') + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--color-jitter', type=float, default=0.4) + parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1') + parser.add_argument('--train-interpolation', type=str, default='bicubic') + parser.add_argument('--reprob', type=float, default=0.25) + parser.add_argument('--remode', type=str, default='pixel') + parser.add_argument('--recount', type=int, default=1) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") + parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging") + parser.add_argument("--wandb-group", default=None, help="wandb group for logging") + parser.add_argument("--wandb-project", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-id", default=None, help="wandb project name for logging") + parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/wandb/run-20251226_094108-svj9s262/files/config.yaml b/wandb/run-20251226_094108-svj9s262/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..46b55aa8284f505239f07f978e8e690c9e314d5b --- /dev/null +++ b/wandb/run-20251226_094108-svj9s262/files/config.yaml @@ -0,0 +1,115 @@ +_wandb: + value: + cli_version: 0.23.0 + code_path: code/src/imagenet/train_model.py + e: + 2v947vqm1whvumuo0j7b7to98dve2h9k: + args: + - --input-size + - "32" + - --data-set + - CIFAR + - --patch-size + - "4" + - --hidden-size + - "128" + - --num-hidden-layers + - "6" + - --warmup-epochs + - "5" + - --num-attention-heads + - "4" + - --intermediate-size + - "512" + - --position-embeddings + - learnable + - --num-labels + - "100" + - --lr + - "1e-4" + - --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 + - /mnt/d/NCKH/weights/lmc/cifar100 + - --data-path + - /mnt/d/NCKH/datasets/cifar100 + codePath: src/imagenet/train_model.py + codePathLocal: src/imagenet/train_model.py + cpu_count: 24 + cpu_count_logical: 24 + cudaVersion: "13.0" + disk: + /: + total: "1081101176832" + used: "67572064256" + email: khanhvinhbui0512@gmail.com + executable: /home/khanhvinh9/miniconda3/envs/lmc/bin/python + git: + commit: 503490314d9f0bef18a30ec7957b3949bd36c751 + remote: https://github.com/MLResearchX/lmc-transformer.git + gpu: NVIDIA GeForce RTX 5060 Laptop GPU + gpu_count: 1 + gpu_nvidia: + - architecture: Blackwell + memoryTotal: "8546942976" + name: NVIDIA GeForce RTX 5060 Laptop GPU + uuid: GPU-6363d584-7737-137f-cc48-e39ef1a9efbd + host: DESKTOP-UCJ5IKF + memory: + total: "16466989056" + os: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 + program: /home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py + python: CPython 3.10.19 + root: /home/khanhvinh9/lmc-transformer + startedAt: "2025-12-26T02:41:08.421096Z" + writerId: 2v947vqm1whvumuo0j7b7to98dve2h9k + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "2": + - 1 + - 5 + - 11 + - 12 + - 41 + - 45 + - 49 + - 51 + - 53 + - 63 + "3": + - 13 + "4": 3.10.19 + "5": 0.23.0 + "6": 4.56.1 + "8": + - 2 + "12": 0.23.0 + "13": linux-x86_64 diff --git a/wandb/run-20251226_094108-svj9s262/files/diff.patch b/wandb/run-20251226_094108-svj9s262/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..64375f8ecec60cacb0096f58826df1afb889a5ae --- /dev/null +++ b/wandb/run-20251226_094108-svj9s262/files/diff.patch @@ -0,0 +1,214 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..ffbb4f8 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3f8f560 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +@@ -78,33 +79,48 @@ def build_dataset(is_train, args): + + def build_transform(is_train, args): + resize_im = args.input_size > 32 +- if is_train: +- # this should always dispatch to transforms_imagenet_train +- transform = create_transform( +- input_size=args.input_size, +- is_training=True, +- color_jitter=args.color_jitter, +- auto_augment=args.aa, +- interpolation=args.train_interpolation, +- re_prob=args.reprob, +- re_mode=args.remode, +- re_count=args.recount, +- ) +- if not resize_im: +- # replace RandomResizedCropAndInterpolation with +- # RandomCrop +- transform.transforms[0] = transforms.RandomCrop( +- args.input_size, padding=4) +- return transform +- t = [] +- if resize_im: +- size = int((256 / 224) * args.input_size) +- t.append( +- transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images +- ) +- t.append(transforms.CenterCrop(args.input_size)) +- t.append(transforms.ToTensor()) +- t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) +- return transforms.Compose(t) ++ if args.data_set == 'CIFAR': ++ if is_train: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.RandomHorizontalFlip(), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ if is_train: ++ # this should always dispatch to transforms_imagenet_train ++ transform = create_transform( ++ input_size=args.input_size, ++ is_training=True, ++ color_jitter=args.color_jitter, ++ auto_augment=args.aa, ++ interpolation=args.train_interpolation, ++ re_prob=args.reprob, ++ re_mode=args.remode, ++ re_count=args.recount, ++ ) ++ if not resize_im: ++ # replace RandomResizedCropAndInterpolation with ++ # RandomCrop ++ transform.transforms[0] = transforms.RandomCrop( ++ args.input_size, padding=4) ++ return transform ++ t = [] ++ if resize_im: ++ size = int((256 / 224) * args.input_size) ++ t.append( ++ transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images ++ ) ++ t.append(transforms.CenterCrop(args.input_size)) ++ t.append(transforms.ToTensor()) ++ t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) ++ return transforms.Compose(t) + + +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251226_094108-svj9s262/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch b/wandb/run-20251226_094108-svj9s262/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch new file mode 100644 index 0000000000000000000000000000000000000000..64375f8ecec60cacb0096f58826df1afb889a5ae --- /dev/null +++ b/wandb/run-20251226_094108-svj9s262/files/diff_503490314d9f0bef18a30ec7957b3949bd36c751.patch @@ -0,0 +1,214 @@ +diff --git a/scripts/imagenet/train_model.sh b/scripts/imagenet/train_model.sh +index ca6dd07..9b60e59 100755 +--- a/scripts/imagenet/train_model.sh ++++ b/scripts/imagenet/train_model.sh +@@ -1,4 +1,3 @@ +-<<<<<<< HEAD + WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ + --num-shared-experts 0 --num-routed-experts 4 --topk 4\ +@@ -16,16 +15,6 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=3 python src/imagenet/train_model.py \ + --num-shared-experts 1 --num-routed-experts 0 --topk 0\ + --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ + --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +-======= +- +-WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ +- --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ +- --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +- --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +- --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut" --wandb-id 0hxu0bbi\ +- --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ +- --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256-shared1-routed0-topk0/last_1291032 +->>>>>>> master + + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=2,3 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 \ +@@ -34,17 +23,8 @@ WANDB_MODE=online CUDA_VISIBLE_DEVICES=4 python src/imagenet/train_model.py \ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet\ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-sinusoidal-epochs300-batch256/last_1170936 + +-<<<<<<< HEAD + # WANDB_MODE=online CUDA_VISIBLE_DEVICES=4,5 python src/imagenet/train_model.py \ + # --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "rope"\ + # --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-MoE" --wandb-entity "vinh-bui0512-hcmut"\ + # --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet \ + # --restore-checkpoint-path /mnt/data/vinhbk/weights/imagenet/lr0.0005-rope-epochs300-batch256/last_300240 +-======= +-# WANDB_MODE=offline CUDA_VISIBLE_DEVICES=6,7 python src/imagenet/train_model.py \ +-# --lr 0.0005 --epochs 300 --batch-size 256 --seed 0 --position-embeddings "learnable"\ +-# --hidden-size 768 --num-hidden-layers 12 --num-attention-heads 12 --intermediate-size 3072 \ +-# --num-shared-experts 1 --num-routed-experts 0 --topk 0\ +-# --wandb-project "LMC-Attention" --wandb-group "ViT-ImageNet-FFN" --wandb-entity "vinh-bui0512-hcmut"\ +-# --save-dir /mnt/data/vinhbk/weights/imagenet --data-path /mnt/data/vinhbk/datasets/imagenet +->>>>>>> master +diff --git a/src/imagenet/__pycache__/datasets.cpython-310.pyc b/src/imagenet/__pycache__/datasets.cpython-310.pyc +index 0780c71..ffbb4f8 100644 +Binary files a/src/imagenet/__pycache__/datasets.cpython-310.pyc and b/src/imagenet/__pycache__/datasets.cpython-310.pyc differ +diff --git a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc +index 551e4a3..cd67dbf 100644 +Binary files a/src/imagenet/__pycache__/lmc_model.cpython-310.pyc and b/src/imagenet/__pycache__/lmc_model.cpython-310.pyc differ +diff --git a/src/imagenet/datasets.py b/src/imagenet/datasets.py +index 85d3f06..3f8f560 100755 +--- a/src/imagenet/datasets.py ++++ b/src/imagenet/datasets.py +@@ -57,13 +57,14 @@ class INatDataset(ImageFolder): + self.samples.append((path_current, target_current_true)) + # __getitem__ and __len__ inherited from ImageFolder + +-def build_dataset(is_train, args): ++def build_dataset(is_train, args,is_generalization=False): + transform = build_transform(is_train, args) + if args.data_set == 'CIFAR': +- dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform) ++ dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform,download=True) + nb_classes = 100 + elif args.data_set == 'IMNET': +- root = os.path.join(args.data_path, 'train' if is_train else 'val') ++ if(is_generalization): root = args.data_path ++ else: root = os.path.join(args.data_path, 'train' if is_train else 'val') + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 1000 + elif args.data_set == 'INAT': +@@ -78,33 +79,48 @@ def build_dataset(is_train, args): + + def build_transform(is_train, args): + resize_im = args.input_size > 32 +- if is_train: +- # this should always dispatch to transforms_imagenet_train +- transform = create_transform( +- input_size=args.input_size, +- is_training=True, +- color_jitter=args.color_jitter, +- auto_augment=args.aa, +- interpolation=args.train_interpolation, +- re_prob=args.reprob, +- re_mode=args.remode, +- re_count=args.recount, +- ) +- if not resize_im: +- # replace RandomResizedCropAndInterpolation with +- # RandomCrop +- transform.transforms[0] = transforms.RandomCrop( +- args.input_size, padding=4) +- return transform +- t = [] +- if resize_im: +- size = int((256 / 224) * args.input_size) +- t.append( +- transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images +- ) +- t.append(transforms.CenterCrop(args.input_size)) +- t.append(transforms.ToTensor()) +- t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) +- return transforms.Compose(t) ++ if args.data_set == 'CIFAR': ++ if is_train: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.RandomHorizontalFlip(), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ return transforms.Compose([ ++ transforms.RandomCrop(args.input_size, padding=4), ++ transforms.ToTensor(), ++ transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), ++ ]) ++ else: ++ if is_train: ++ # this should always dispatch to transforms_imagenet_train ++ transform = create_transform( ++ input_size=args.input_size, ++ is_training=True, ++ color_jitter=args.color_jitter, ++ auto_augment=args.aa, ++ interpolation=args.train_interpolation, ++ re_prob=args.reprob, ++ re_mode=args.remode, ++ re_count=args.recount, ++ ) ++ if not resize_im: ++ # replace RandomResizedCropAndInterpolation with ++ # RandomCrop ++ transform.transforms[0] = transforms.RandomCrop( ++ args.input_size, padding=4) ++ return transform ++ t = [] ++ if resize_im: ++ size = int((256 / 224) * args.input_size) ++ t.append( ++ transforms.Resize(size, interpolation=3), # to maintain same ratio w.r.t. 224 images ++ ) ++ t.append(transforms.CenterCrop(args.input_size)) ++ t.append(transforms.ToTensor()) ++ t.append(transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))) ++ return transforms.Compose(t) + + +diff --git a/src/imagenet/matching_utils.py b/src/imagenet/matching_utils.py +index b913c42..4f47b19 100644 +--- a/src/imagenet/matching_utils.py ++++ b/src/imagenet/matching_utils.py +@@ -446,7 +446,7 @@ def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activa + def weight_matching_attn(rng, params_a, params_b, activation, config): + params_dict = {} + configurations = [ +- ("permu_head_init_ortho_no_opt", 'ortho', True, False), ++ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), + ("permu_head_init_ortho_opt", 'ortho', True, True), + # ("naive_head_init_ortho_no_opt", 'ortho', False, False), + # ("naive_head_init_ortho_opt", 'ortho', False, True), +diff --git a/src/imagenet/train_model.py b/src/imagenet/train_model.py +index 10b031a..c48bfef 100755 +--- a/src/imagenet/train_model.py ++++ b/src/imagenet/train_model.py +@@ -25,7 +25,7 @@ from pprint import pprint + from typing import Any, Dict, List + import shutil + mp.set_start_method("spawn", force=True) +-os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7" ++os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9" + def aggregate_metrics(metrics_list): + return { + key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item() +@@ -74,11 +74,11 @@ def main(args: argparse.Namespace): + entity=args.wandb_entity, + group=args.wandb_group, + id=args.wandb_id, +- name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}", ++ name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}", + save_code=True + ) + wandb.config = dict(vars(args)) +- save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-routed{args.num_routed_experts}-topk{args.topk}") ++ save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}") + torch.manual_seed(args.seed) + np.random.seed(args.seed) + torch.cuda.manual_seed_all(args.seed) +@@ -86,11 +86,14 @@ def main(args: argparse.Namespace): + # --- Prepare Data Loader --- + train_loader, val_loader = imagenet_data_loader(args) + # --- Load pretrained model --- +- config = ViTConfig.from_pretrained('google/vit-base-patch16-224') ++ config = ViTConfig() + config.hidden_size = args.hidden_size + config.num_hidden_layers = args.num_hidden_layers + config.num_attention_heads = args.num_attention_heads + config.intermediate_size = args.intermediate_size ++ config.patch_size = args.patch_size ++ config.image_size = args.input_size ++ config.num_labels = args.num_labels + config.position_embeddings = args.position_embeddings + config.rotary_value = args.rotary_value + config.num_shared_experts = args.num_shared_experts +@@ -210,6 +213,8 @@ if __name__ == "__main__": + parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.") + parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.") + parser.add_argument("--position-embeddings", type=str, default='sinusoidal') ++ parser.add_argument("--patch-size", type=int, default=16, help="") ++ parser.add_argument("--num-labels", type=int, default=1000, help="") + parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.') + parser.add_argument("--num-shared-experts", type=int, default = 1) + parser.add_argument("--num-routed-experts", type=int, default = 0) diff --git a/wandb/run-20251226_094108-svj9s262/files/output.log b/wandb/run-20251226_094108-svj9s262/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..9120a4f462545476f862f8673b10b02e06f6e060 --- /dev/null +++ b/wandb/run-20251226_094108-svj9s262/files/output.log @@ -0,0 +1,53 @@ +TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We recommend migrating to PyTorch classes or pinning your version of Transformers. +Starting training... +JAX devices: [CudaDevice(id=0)] +Using 1 devices +Epoch 0 +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 166, in main + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 494, in __iter__ + return self._get_iterator() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 427, in _get_iterator + return _MultiProcessingDataLoaderIter(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1170, in __init__ + w.start() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/process.py", line 121, in start + self._popen = self._Popen(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 224, in _Popen + return _default_context.get_context().Process._Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 288, in _Popen + return Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__ + super().__init__(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__ + self._launch(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 62, in _launch + f.write(fp.getbuffer()) +KeyboardInterrupt +Traceback (most recent call last): + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 264, in + main(args) + File "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", line 166, in main + pbar = tqdm(enumerate(train_loader), desc="Training", leave=False) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 494, in __iter__ + return self._get_iterator() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 427, in _get_iterator + return _MultiProcessingDataLoaderIter(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1170, in __init__ + w.start() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/process.py", line 121, in start + self._popen = self._Popen(self) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 224, in _Popen + return _default_context.get_context().Process._Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/context.py", line 288, in _Popen + return Popen(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__ + super().__init__(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__ + self._launch(process_obj) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 62, in _launch + f.write(fp.getbuffer()) +KeyboardInterrupt diff --git a/wandb/run-20251226_094108-svj9s262/files/requirements.txt b/wandb/run-20251226_094108-svj9s262/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a1d05b9c95dd298711744d872e6fda72ccd954e --- /dev/null +++ b/wandb/run-20251226_094108-svj9s262/files/requirements.txt @@ -0,0 +1,140 @@ +frozenlist==1.8.0 +idna==3.11 +joblib==1.5.2 +yarl==1.22.0 +charset-normalizer==3.4.4 +httpx==0.28.1 +smmap==5.0.2 +nvidia-cuda-nvrtc-cu12==12.8.93 +nest-asyncio==1.6.0 +protobuf==6.33.1 +hf-xet==1.2.0 +sentry-sdk==2.44.0 +watchfiles==1.1.1 +opt_einsum==3.4.0 +toml==0.10.2 +types-certifi==2021.10.8.3 +hyperframe==6.1.0 +six==1.17.0 +aiohappyeyeballs==2.6.1 +grpclib==0.4.8 +pip==25.3 +nvidia-cufile-cu12==1.13.1.3 +datasets==4.4.1 +nvidia-cublas-cu12==12.8.4.1 +nvidia-curand-cu12==10.3.9.90 +GitPython==3.1.45 +timm==1.0.22 +shellingham==1.5.4 +pyparsing==3.2.5 +safetensors==0.6.2 +orbax-checkpoint==0.11.28 +tqdm==4.67.1 +annotated-types==0.7.0 +platformdirs==4.5.0 +aiofiles==25.1.0 +aiosignal==1.4.0 +nvidia-cudnn-cu12==9.10.2.21 +dill==0.4.0 +nvidia-cusparse-cu12==12.5.8.93 +cbor2==5.7.1 +pytz==2025.2 +psutil==7.1.3 +attrs==25.4.0 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvtx-cu12==12.8.90 +nvidia-cusparselt-cu12==0.7.1 +urllib3==2.5.0 +kiwisolver==1.4.9 +typer-slim==0.20.0 +mdurl==0.1.2 +sympy==1.14.0 +pydantic==2.11.10 +importlib_resources==6.5.2 +h11==0.16.0 +wheel==0.45.1 +nvitop==1.6.0 +msgpack==1.1.2 +pandas==2.3.3 +markdown-it-py==4.0.0 +certifi==2025.11.12 +multiprocess==0.70.18 +async-timeout==5.0.1 +nvidia-cuda-cupti-cu12==12.8.90 +propcache==0.4.1 +threadpoolctl==3.6.0 +exceptiongroup==1.3.0 +chex==0.1.90 +matplotlib==3.10.7 +pillow==12.0.0 +typing_extensions==4.15.0 +toolz==1.1.0 +httpcore==1.0.9 +tzdata==2025.2 +jax-cuda12-pjrt==0.6.2 +typing-inspection==0.4.2 +pyarrow==22.0.0 +python-dateutil==2.9.0.post0 +requests==2.32.5 +networkx==3.4.2 +anyio==4.11.0 +nvidia-cuda-nvcc-cu12==12.9.86 +mpmath==1.3.0 +tokenizers==0.22.1 +nvidia-ml-py==13.580.82 +MarkupSafe==3.0.3 +Pygments==2.19.2 +cycler==0.12.1 +zipp==3.23.0 +jaxlib==0.6.2 +typer==0.20.0 +PyYAML==6.0.3 +scikit-learn==1.7.2 +argparse==1.4.0 +nvidia-cuda-runtime-cu12==12.8.90 +etils==1.13.0 +h2==4.3.0 +nvidia-nvshmem-cu12==3.3.20 +triton==3.5.1 +regex==2025.11.3 +rich==14.2.0 +nvidia-cufft-cu12==11.3.3.83 +treescope==0.1.10 +contourpy==1.3.2 +xxhash==3.6.0 +wandb==0.23.0 +jax-cuda12-plugin==0.6.2 +simplejson==3.20.2 +dataclasses==0.6 +jax==0.6.2 +fonttools==4.60.1 +huggingface-hub==0.36.0 +torch==2.9.1 +humanize==4.14.0 +nvidia-cusolver-cu12==11.7.3.90 +multidict==6.7.0 +types-toml==0.10.8.20240310 +hpack==4.1.0 +optax==0.2.6 +fsspec==2025.10.0 +filelock==3.20.0 +synchronicity==0.10.4 +tensorstore==0.1.78 +transformers==4.56.1 +ml_dtypes==0.5.3 +pydantic_core==2.33.2 +gitdb==4.0.12 +torchvision==0.24.1 +setuptools==80.9.0 +aiohttp==3.13.2 +sniffio==1.3.1 +scipy==1.15.3 +click==8.3.1 +Jinja2==3.1.6 +modal==1.2.4 +numpy==2.2.6 +packaging==25.0 +nvidia-nccl-cu12==2.27.5 +absl-py==2.3.1 +flax==0.10.7 +LMC-of-MoE==0.3.1 diff --git a/wandb/run-20251226_094108-svj9s262/files/wandb-metadata.json b/wandb/run-20251226_094108-svj9s262/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..b9d6f4fc75e53962b0f67337077aaaeadd793c3d --- /dev/null +++ b/wandb/run-20251226_094108-svj9s262/files/wandb-metadata.json @@ -0,0 +1,85 @@ +{ + "os": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "python": "CPython 3.10.19", + "startedAt": "2025-12-26T02:41:08.421096Z", + "args": [ + "--input-size", + "32", + "--data-set", + "CIFAR", + "--patch-size", + "4", + "--hidden-size", + "128", + "--num-hidden-layers", + "6", + "--warmup-epochs", + "5", + "--num-attention-heads", + "4", + "--intermediate-size", + "512", + "--position-embeddings", + "learnable", + "--num-labels", + "100", + "--lr", + "1e-4", + "--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", + "/mnt/d/NCKH/weights/lmc/cifar100", + "--data-path", + "/mnt/d/NCKH/datasets/cifar100" + ], + "program": "/home/khanhvinh9/lmc-transformer/src/imagenet/train_model.py", + "codePath": "src/imagenet/train_model.py", + "codePathLocal": "src/imagenet/train_model.py", + "git": { + "remote": "https://github.com/MLResearchX/lmc-transformer.git", + "commit": "503490314d9f0bef18a30ec7957b3949bd36c751" + }, + "email": "khanhvinhbui0512@gmail.com", + "root": "/home/khanhvinh9/lmc-transformer", + "host": "DESKTOP-UCJ5IKF", + "executable": "/home/khanhvinh9/miniconda3/envs/lmc/bin/python", + "cpu_count": 24, + "cpu_count_logical": 24, + "gpu": "NVIDIA GeForce RTX 5060 Laptop GPU", + "gpu_count": 1, + "disk": { + "/": { + "total": "1081101176832", + "used": "67572064256" + } + }, + "memory": { + "total": "16466989056" + }, + "gpu_nvidia": [ + { + "name": "NVIDIA GeForce RTX 5060 Laptop GPU", + "memoryTotal": "8546942976", + "architecture": "Blackwell", + "uuid": "GPU-6363d584-7737-137f-cc48-e39ef1a9efbd" + } + ], + "cudaVersion": "13.0", + "writerId": "2v947vqm1whvumuo0j7b7to98dve2h9k" +} \ No newline at end of file diff --git a/wandb/run-20251226_094108-svj9s262/files/wandb-summary.json b/wandb/run-20251226_094108-svj9s262/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..26eb9fe6450816c7c6ef2f464335d313076bb440 --- /dev/null +++ b/wandb/run-20251226_094108-svj9s262/files/wandb-summary.json @@ -0,0 +1 @@ +{"_wandb":{"runtime":19},"_runtime":19} \ No newline at end of file diff --git a/wandb/run-20251226_094108-svj9s262/logs/debug-core.log b/wandb/run-20251226_094108-svj9s262/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..decbab02e87233740f432ab8ae88249740b32635 --- /dev/null +++ b/wandb/run-20251226_094108-svj9s262/logs/debug-core.log @@ -0,0 +1,12 @@ +{"time":"2025-12-26T09:41:08.673264346+07:00","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmp4hdqi1u1/port-1318853.txt","pid":1318853,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false} +{"time":"2025-12-26T09:41:08.675707349+07:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":1318853} +{"time":"2025-12-26T09:41:08.675705698+07:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-1318853-1319015-1883696914/socket","Net":"unix"}} +{"time":"2025-12-26T09:41:08.760580682+07:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"} +{"time":"2025-12-26T09:41:08.772990303+07:00","level":"INFO","msg":"handleInformInit: received","streamId":"svj9s262","id":"1(@)"} +{"time":"2025-12-26T09:41:09.250020354+07:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"svj9s262","id":"1(@)"} +{"time":"2025-12-26T09:41:29.560138295+07:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"} +{"time":"2025-12-26T09:41:29.560241626+07:00","level":"INFO","msg":"connection: closing","id":"1(@)"} +{"time":"2025-12-26T09:41:29.560259582+07:00","level":"INFO","msg":"connection: closed successfully","id":"1(@)"} +{"time":"2025-12-26T09:41:29.560262008+07:00","level":"INFO","msg":"server is shutting down"} +{"time":"2025-12-26T09:41:29.560296802+07:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-1318853-1319015-1883696914/socket","Net":"unix"}} +{"time":"2025-12-26T09:41:30.357212149+07:00","level":"INFO","msg":"server: parent process exited, terminating service process"} diff --git a/wandb/run-20251226_094108-svj9s262/logs/debug-internal.log b/wandb/run-20251226_094108-svj9s262/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..0658c12996a5f0c436205033c2de8f24e45ef307 --- /dev/null +++ b/wandb/run-20251226_094108-svj9s262/logs/debug-internal.log @@ -0,0 +1,7 @@ +{"time":"2025-12-26T09:41:08.773359688+07:00","level":"INFO","msg":"stream: starting","core version":"0.23.0"} +{"time":"2025-12-26T09:41:09.249853643+07:00","level":"INFO","msg":"stream: created new stream","id":"svj9s262"} +{"time":"2025-12-26T09:41:09.249933283+07:00","level":"INFO","msg":"handler: started","stream_id":"svj9s262"} +{"time":"2025-12-26T09:41:09.250011964+07:00","level":"INFO","msg":"stream: started","id":"svj9s262"} +{"time":"2025-12-26T09:41:09.250090029+07:00","level":"INFO","msg":"writer: started","stream_id":"svj9s262"} +{"time":"2025-12-26T09:41:09.250121649+07:00","level":"INFO","msg":"sender: started","stream_id":"svj9s262"} +{"time":"2025-12-26T09:41:29.560166541+07:00","level":"INFO","msg":"stream: closing","id":"svj9s262"} diff --git a/wandb/run-20251226_094108-svj9s262/logs/debug.log b/wandb/run-20251226_094108-svj9s262/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..1b4ceac7a0a901622a983c980755b3d7002fd51d --- /dev/null +++ b/wandb/run-20251226_094108-svj9s262/logs/debug.log @@ -0,0 +1,348 @@ +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_setup.py:_flush():80] Current SDK version is 0.23.0 +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_setup.py:_flush():80] Configure stats pid to 1318853 +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/.config/wandb/settings +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_setup.py:_flush():80] Loading settings from /home/khanhvinh9/lmc-transformer/wandb/settings +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_setup.py:_flush():80] Loading settings from environment variables +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_init.py:setup_run_log_directory():713] Logging user logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251226_094108-svj9s262/logs/debug.log +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_init.py:setup_run_log_directory():714] Logging internal logs to /home/khanhvinh9/lmc-transformer/wandb/run-20251226_094108-svj9s262/logs/debug-internal.log +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_init.py:init():840] calling init triggers +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_init.py:init():845] wandb.init called with sweep_config: {} +config: {'_wandb': {'code_path': 'code/src/imagenet/train_model.py'}} +2025-12-26 09:41:08,427 INFO MainThread:1318853 [wandb_init.py:init():888] starting backend +2025-12-26 09:41:08,759 INFO MainThread:1318853 [wandb_init.py:init():891] sending inform_init request +2025-12-26 09:41:08,766 INFO MainThread:1318853 [wandb_init.py:init():899] backend started and connected +2025-12-26 09:41:08,768 INFO MainThread:1318853 [wandb_init.py:init():969] updated telemetry +2025-12-26 09:41:08,772 INFO MainThread:1318853 [wandb_init.py:init():993] communicating run to backend with 90.0 second timeout +2025-12-26 09:41:09,748 INFO MainThread:1318853 [wandb_init.py:init():1040] starting run threads in backend +2025-12-26 09:41:09,895 INFO MainThread:1318853 [wandb_run.py:_console_start():2504] atexit reg +2025-12-26 09:41:09,895 INFO MainThread:1318853 [wandb_run.py:_redirect():2352] redirect: wrap_raw +2025-12-26 09:41:09,895 INFO MainThread:1318853 [wandb_run.py:_redirect():2421] Wrapping output streams. +2025-12-26 09:41:09,895 INFO MainThread:1318853 [wandb_run.py:_redirect():2444] Redirects installed. +2025-12-26 09:41:09,898 INFO MainThread:1318853 [wandb_init.py:init():1080] run started, returning control to user process +2025-12-26 09:41:29,560 INFO wandb-AsyncioManager-main:1318853 [service_client.py:_forward_responses():80] Reached EOF. +2025-12-26 09:41:29,560 INFO wandb-AsyncioManager-main:1318853 [mailbox.py:close():137] Closing mailbox, abandoning 1 handles. +2025-12-26 09:41:29,931 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:29,938 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:29,939 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:29,944 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:29,975 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,011 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,022 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,091 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,108 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,108 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,109 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,110 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,112 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,112 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,112 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,117 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,118 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,118 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,119 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost +2025-12-26 09:41:30,119 ERROR wandb-AsyncioManager-main:1318853 [asyncio_manager.py:fn_wrap_exceptions():183] Uncaught exception in run_soon callback. +Traceback (most recent call last): + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/asyncio_manager.py", line 181, in fn_wrap_exceptions + await fn() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 38, in publish + await self._send_server_request(request) + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/site-packages/wandb/sdk/lib/service/service_client.py", line 64, in _send_server_request + await self._writer.drain() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 371, in drain + await self._protocol._drain_helper() + File "/home/khanhvinh9/miniconda3/envs/lmc/lib/python3.10/asyncio/streams.py", line 167, in _drain_helper + raise ConnectionResetError('Connection lost') +ConnectionResetError: Connection lost diff --git a/wandb/run-20251226_094108-svj9s262/run-svj9s262.wandb b/wandb/run-20251226_094108-svj9s262/run-svj9s262.wandb new file mode 100644 index 0000000000000000000000000000000000000000..30977402b08eb894463b83d7de7a007255879a32 Binary files /dev/null and b/wandb/run-20251226_094108-svj9s262/run-svj9s262.wandb differ