{"id":22495435,"url":"https://github.com/saforem2/mmm","last_synced_at":"2025-08-13T01:05:28.022Z","repository":{"id":266412273,"uuid":"897940794","full_name":"saforem2/mmm","owner":"saforem2","description":"Multi-Modal Modeling","archived":false,"fork":false,"pushed_at":"2025-06-16T12:59:02.000Z","size":371,"stargazers_count":6,"open_issues_count":2,"forks_count":0,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-08-07T22:34:52.994Z","etag":null,"topics":["distributed-training","llm","multi-modal","pytorch"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"bsd-3-clause","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/saforem2.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2024-12-03T14:05:48.000Z","updated_at":"2025-05-24T05:17:26.000Z","dependencies_parsed_at":"2024-12-04T06:27:14.355Z","dependency_job_id":"c87aee05-1a92-495e-870f-f539376c6231","html_url":"https://github.com/saforem2/mmm","commit_stats":null,"previous_names":["saforem2/mmm"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/saforem2/mmm","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/saforem2%2Fmmm","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/saforem2%2Fmmm/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/saforem2%2Fmmm/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/saforem2%2Fmmm/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/saforem2","download_url":"https://codeload.github.com/saforem2/mmm/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/saforem2%2Fmmm/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":270160883,"owners_count":24537534,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-08-12T02:00:09.011Z","response_time":80,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["distributed-training","llm","multi-modal","pytorch"],"created_at":"2024-12-06T19:28:17.652Z","updated_at":"2025-08-13T01:05:27.998Z","avatar_url":"https://github.com/saforem2.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# 🐫 `mmm`: Multi-Modal Modeling\n\n## 🐣 Getting Started\n\n### 🏡 Setup Environment\n\n\u003e [!IMPORTANT]\n\u003e See 🍋[saforem2 / `ezpz`](https://github.com/saforem2/ezpz) for additional information.\n\nWe use [`ezpz`](https://github.com/saforem2/ezpz)\nfor **setting up**, **launching**, and **orchestrating** our distributed training.\n\nTo setup our environment, we first `source` the\n[`ezpz/bin/utils.sh`](https://github.com/saforem2/ezpz/blob/main/src/ezpz/bin/utils.sh)\nscript.\n\nThis populates the shell environment with a variety of useful functions for\nsetting up a python environment and determining the specifics of \nour currently running job (e.g. number of nodes, number of GPUs per node, etc.).\n\nIn particular, we can use the `ezpz_setup_env` helper function[^ezpz_setup_env]\nto automatically take care of all of our required environment setup:\n\n```bash\n# if @ ALCF, set:\n# PBS_O_WORKDIR=$(pwd)\nsource \u003c(curl -s https://raw.githubusercontent.com/saforem2/ezpz/refs/heads/main/src/ezpz/bin/utils.sh)\nezpz_setup_env\n```\n\n- 🪄 This will, _automagically_:\n\n    - 🐍 Setup + activate python environment\n    - 🤔 Determine available resources (i.e. `NHOSTS`, `NGPU_PER_HOST`, `NGPUS`)\n    - 🚀 Define a `launch` alias to launch our application across them\n \n    and should work anywhere you have a working MPI installation.\n\n\u003e [!NOTE]\n\u003e This is _technically_ optional, but highly recommended as it will allow\n\u003e you to automatically launch on any[^any] distributed setup with a compatible MPI.\n\n[^ezpz_setup_env]:\n    Technically, it just chains together two separate (and useful on their own) function calls, explicitly:\n\n    ```bash\n    $ which ezpz_setup_env\n    ezpz_setup_env() {\n       ezpz_setup_python \u0026\u0026 ezpz_setup_job\n    }\n    ```\n\n\n[^any]: This has been tested and confirmed to work on:\n\n    - Any job behind a {PBS, slurm} job scheduler\n    - All [ALCF](https://alcf.anl.gov) systems (Aurora, Polaris, Sunspot, Sophia, etc.)\n    - Frontier (AMD system) @ OLCF\n    - Perlmutter (NVIDIA system) @ NERSC\n    - Distributed CPUs via `mpirun`\n    - Distributed `mps` devices via `mpirun`\n\n    Both PBS and Slurm job schedulers are supported and the specifics of the running job will be used to populate the corresponding `launch` command.\n\n### ⬇️ Install\n\nArmed with a functional python installation[^requires], we can install `mmm`.\n\n[^requires]: \n    Requires:\n    \n    - 🔥 [pytorch/`pytorch`](https://pytorch.org)\n    - 🍋 [saforem2/`ezpz`](https://github.com/saforem2/ezpz)\n    - 📡 [mpi4py/`mpi4py`](https://github.com/mpi4py/mpi4py)\n\n\n- From GitHub:\n\n  ```bash\n  python3 -m pip install -e \"git+https://github.com/saforem2/mmm#egg=mmm\"\n  ```\n\n- \u003cdetails closed\u003e\u003csummary\u003eFrom local clone (development)\u003c/summary\u003e\n\n  ```bash\n  git clone https://github.com/saforem2/mmm\n  python3 -m pip install -e mmm\n  ```\n\n  \u003c/details\u003e\n\n- \u003cdetails closed\u003e\u003csummary\u003eUsing \u003ccode\u003euv\u003c/code\u003e\u003c/summary\u003e\n\n  ```bash\n  uv pip install \"mmm @ git+https://github.com/saforem2/mmm\"\n  ```\n\n  \u003c/details\u003e\n\n## 📊 Examples\n\n### 🕸️ TP + FSDP on Aurora\n\n- See [`src/mmm/examples/fsdp_tp.py`](/src/mmm/examples/fsdp_tp.py)\n\n- We use a combination of:\n\n  1. Tensor Parallelism (TP)\n  1. Fully Sharded Data Parallelism (FSDP) (across data-parallel (DP) groups)\n\n- Launch:\n\n    ```bash\n    launch python3 -m mmm.examples.fsdp_tp --n_layers 24 --tpsize 4\n    ```\n\n    \u003cdetails closed\u003e\u003csummary\u003eOutput:\u003c/summary\u003e\n\n\n    \u003cdetails closed\u003e\u003csummary\u003eAurora @ ALCF:\u003c/summary\u003e\n\n    ```python\n    # $ launch python3 -m mmm.examples.fsdp_tp --tp 4 --epochs 5 --batch-size 4 --dim=4096 | egrep -vi 'cu|CCL'\n    [2025-02-02 15:25:42][I][datasets/config:54:datasets] PyTorch version 2.3.1+cxx11.abi available.\n    [2025-02-02 15:25:42][I][datasets/config:112:datasets] TensorFlow version 2.15.1 available.\n    [2025-02-02 15:25:42][I][datasets/config:125:datasets] JAX version 0.5.0 available.\n    [2025-02-02 15:25:42][I][ezpz/dist:505] Using get_torch_device_type()='xpu' with backend='ccl'\n    [2025-02-02 15:25:42][I][ezpz/dist:881] ['x4515c0s1b0n0'][ 2/23]\n    [2025-02-02 15:25:42][I][ezpz/dist:881] ['x4515c0s1b0n0'][ 1/23]\n    [2025-02-02 15:25:42][I][ezpz/dist:881] ['x4515c0s1b0n0'][ 3/23]\n    [2025-02-02 15:25:42][I][ezpz/dist:881] ['x4515c0s1b0n0'][ 5/23]\n    [2025-02-02 15:25:42][I][ezpz/dist:881] ['x4515c0s1b0n0'][ 4/23]\n    [2025-02-02 15:25:42][I][ezpz/dist:881] ['x4515c0s1b0n0'][ 9/23]\n    [2025-02-02 15:25:42][I][ezpz/dist:881] ['x4515c0s1b0n0'][10/23]\n    [2025-02-02 15:25:42][I][ezpz/dist:881] ['x4515c0s1b0n0'][ 8/23]\n    [2025-02-02 15:25:42][I][ezpz/dist:881] ['x4515c0s1b0n0'][11/23]\n    [2025-02-02 15:25:42][I][ezpz/dist:881] ['x4515c0s1b0n0'][ 6/23]\n    [2025-02-02 15:25:42][I][ezpz/dist:881] ['x4515c0s1b0n0'][ 7/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][19/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][15/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][12/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][13/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][16/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][14/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][17/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][18/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][22/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][23/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][21/23]\n    [2025-02-02 15:25:43][I][tp/__init__:146:ezpz.tp] \u003e initializing tensor parallel with size 4\n    [2025-02-02 15:25:43][I][tp/__init__:151:ezpz.tp] \u003e initializing context parallel with size 1\n    [2025-02-02 15:25:43][I][tp/__init__:156:ezpz.tp] \u003e initializing pipeline with size 1\n    [2025-02-02 15:25:43][I][tp/__init__:159:ezpz.tp] \u003e initializing ddp with size 6\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s2b0n0'][20/23]\n    [2025-02-02 15:25:43][I][ezpz/dist:835] Using device='xpu' with backend='DDP' + 'ccl' for distributed training.\n    [2025-02-02 15:25:43][I][ezpz/dist:881] ['x4515c0s1b0n0'][ 0/23]\n    [2025-02-02 15:25:43][I][examples/fsdp_tp:177:__main__] Device mesh created:\n    device_mesh=DeviceMesh([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]], mesh_dim_names=('dp', 'tp'))\n    [2025-02-02 15:25:43][I][examples/fsdp_tp:187:__main__] config:\n    ModelArgs(dim=4096, n_layers=32, n_heads=32, n_kv_heads=4, vocab_size=32000, multiple_of=360, ffn_dim_multiplier=None, norm_eps=1e-05, max_batch_size=32, max_seq_len=32768, depth_init=True)\n    [2025-02-02 15:26:29][I][examples/fsdp_tp:191:__main__]\n    =================================================================\n    Layer (type:depth-idx)                   Param #\n    =================================================================\n    Transformer                              --\n    ├─Embedding: 1-1                         131,072,000\n    ├─ModuleList: 1-2                        --\n    │    └─TransformerBlock: 2-1             174,891,008\n    │    └─TransformerBlock: 2-2             174,891,008\n    │    └─TransformerBlock: 2-3             174,891,008\n    │    └─TransformerBlock: 2-4             174,891,008\n    │    └─TransformerBlock: 2-5             174,891,008\n    │    └─TransformerBlock: 2-6             174,891,008\n    │    └─TransformerBlock: 2-7             174,891,008\n    │    └─TransformerBlock: 2-8             174,891,008\n    │    └─TransformerBlock: 2-9             174,891,008\n    │    └─TransformerBlock: 2-10            174,891,008\n    │    └─TransformerBlock: 2-11            174,891,008\n    │    └─TransformerBlock: 2-12            174,891,008\n    │    └─TransformerBlock: 2-13            174,891,008\n    │    └─TransformerBlock: 2-14            174,891,008\n    │    └─TransformerBlock: 2-15            174,891,008\n    │    └─TransformerBlock: 2-16            174,891,008\n    │    └─TransformerBlock: 2-17            174,891,008\n    │    └─TransformerBlock: 2-18            174,891,008\n    │    └─TransformerBlock: 2-19            174,891,008\n    │    └─TransformerBlock: 2-20            174,891,008\n    │    └─TransformerBlock: 2-21            174,891,008\n    │    └─TransformerBlock: 2-22            174,891,008\n    │    └─TransformerBlock: 2-23            174,891,008\n    │    └─TransformerBlock: 2-24            174,891,008\n    │    └─TransformerBlock: 2-25            174,891,008\n    │    └─TransformerBlock: 2-26            174,891,008\n    │    └─TransformerBlock: 2-27            174,891,008\n    │    └─TransformerBlock: 2-28            174,891,008\n    │    └─TransformerBlock: 2-29            174,891,008\n    │    └─TransformerBlock: 2-30            174,891,008\n    │    └─TransformerBlock: 2-31            174,891,008\n    │    └─TransformerBlock: 2-32            174,891,008\n    ├─RMSNorm: 1-3                           4,096\n    ├─Linear: 1-4                            131,072,000\n    =================================================================\n    Total params: 5,858,660,352\n    Trainable params: 5,858,660,352\n    Non-trainable params: 0\n    =================================================================\n    2025:02:02-15:26:31:( 3698) |CCL_WARN| value of CCL_LOG_LEVEL changed to be error (default:warn)\n    [2025-02-02 15:26:31][I][examples/fsdp_tp:163:__main__] Model after parallelization:\n    sharded_model=FullyShardedDataParallel(\n      (_fsdp_wrapped_module): Transformer(\n        (tok_embeddings): Embedding(32000, 4096)\n        (layers): ModuleList(\n          (0-31): 32 x TransformerBlock(\n            (attention): Attention(\n              (wq): Linear(in_features=4096, out_features=4096, bias=False)\n              (wk): Linear(in_features=4096, out_features=512, bias=False)\n              (wv): Linear(in_features=4096, out_features=512, bias=False)\n              (wo): Linear(in_features=4096, out_features=4096, bias=False)\n            )\n            (feed_forward): FeedForward(\n              (w1): Linear(in_features=4096, out_features=11160, bias=False)\n              (w2): Linear(in_features=11160, out_features=4096, bias=False)\n              (w3): Linear(in_features=4096, out_features=11160, bias=False)\n            )\n            (attention_norm): RMSNorm()\n            (ffn_norm): RMSNorm()\n          )\n        )\n        (norm): RMSNorm()\n        (output): Linear(in_features=4096, out_features=32000, bias=False)\n      )\n    )\n    \n    [2025-02-02 15:26:31][I][examples/fsdp_tp:196:__main__] Creating optimizer=AdamW with lr=5e-05\n    [2025-02-02 15:26:31][I][examples/fsdp_tp:245:__main__] Starting 2D training...\n    [2025-02-02 15:28:55][I][examples/fsdp_tp:281:__main__] epoch=0 iter=0 loss=10.852120 dt=37.205128 dtf=30.174712 dtb=7.030416\n    [2025-02-02 15:28:55][I][examples/fsdp_tp:294:__main__] inp.shape=torch.Size([4, 2048])\n    [2025-02-02 15:28:59][I][examples/fsdp_tp:281:__main__] epoch=0 iter=1 loss=10.844956 dt=4.255044 dtf=0.894840 dtb=3.360204\n    [2025-02-02 15:29:03][I][examples/fsdp_tp:281:__main__] epoch=0 iter=2 loss=10.852772 dt=4.236974 dtf=1.022428 dtb=3.214546\n    [2025-02-02 15:29:08][I][examples/fsdp_tp:281:__main__] epoch=0 iter=3 loss=10.845860 dt=4.306954 dtf=1.042703 dtb=3.264251\n    [2025-02-02 15:29:12][I][examples/fsdp_tp:281:__main__] epoch=0 iter=4 loss=10.847977 dt=4.324451 dtf=1.043026 dtb=3.281424\n    [2025-02-02 15:29:16][I][examples/fsdp_tp:281:__main__] epoch=0 iter=5 loss=10.861067 dt=4.322672 dtf=1.049985 dtb=3.272688\n    [2025-02-02 15:29:21][I][examples/fsdp_tp:281:__main__] epoch=0 iter=6 loss=10.875711 dt=4.268926 dtf=1.038222 dtb=3.230704\n    [2025-02-02 15:29:25][I][examples/fsdp_tp:281:__main__] epoch=0 iter=7 loss=10.871173 dt=4.311684 dtf=1.041232 dtb=3.270452\n    [2025-02-02 15:29:29][I][examples/fsdp_tp:281:__main__] epoch=0 iter=8 loss=10.860594 dt=4.271564 dtf=1.032435 dtb=3.239130\n    [2025-02-02 15:29:34][I][examples/fsdp_tp:281:__main__] epoch=0 iter=9 loss=10.863268 dt=4.330087 dtf=1.058527 dtb=3.271560\n    [2025-02-02 15:29:38][I][examples/fsdp_tp:281:__main__] epoch=0 iter=10 loss=10.865101 dt=4.256579 dtf=1.028928 dtb=3.227651\n    ```\n    \n    \u003c/details\u003e\n    \n    \u003cdetails closed\u003e\u003csummary\u003ePolaris @ ALCF:\u003c/summary\u003e\n    \n    ```python\n    # launch python3 -Wignore -m mmm.train --job.config_file train_configs/debug_model.toml\n    #[01:57:26 PM][x3002c0s25b1n0][/e/a/f/p/s/mmm][🌱 main][🤷✓] [⏱️ 1m14s]\n    ; launch python3 -m mmm.examples.fsdp_tp --tp 4 --epochs 5 --batch-size 4 --dim=4096\n    Connected to tcp://x3002c0s13b0n0.hsn.cm.polaris.alcf.anl.gov:7919\n    Found executable /lus/eagle/projects/argonne_tpc/foremans/projects/saforem2/mmm/venvs/2024-04-29/bin/python3\n    Launching application a1bd5be5-8971-4601-9352-acef885601d6\n    Using PMI port 36863,36864\n    [2025-02-01 13:57:40][I][datasets/config:58:datasets] PyTorch version 2.3.0 available.\n    [2025-02-01 13:57:40][I][datasets/config:105:datasets] TensorFlow version 2.16.1 available.\n    [2025-02-01 13:57:40][I][datasets/config:118:datasets] JAX version 0.4.26 available.\n    [2025-02-01 13:57:40][I][tp/__init__:146:ezpz.tp] \u003e initializing tensor parallel with size 4\n    [2025-02-01 13:57:40][I][tp/__init__:151:ezpz.tp] \u003e initializing context parallel with size 1\n    [2025-02-01 13:57:40][I][tp/__init__:156:ezpz.tp] \u003e initializing pipeline with size 1\n    [2025-02-01 13:57:40][I][tp/__init__:159:ezpz.tp] \u003e initializing ddp with size 8\n    [2025-02-01 13:57:42][I][ezpz/dist:823] Using device='cuda' with backend='DDP' + 'nccl' for distributed training.\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s13b0n0'][ 3/31] [tp:3/3][dp:0/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s13b1n0'][ 7/31] [tp:3/3][dp:1/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s1b0n0'][19/31] [tp:3/3][dp:4/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s19b0n0'][11/31] [tp:3/3][dp:2/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s25b0n0'][27/31] [tp:3/3][dp:6/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s1b1n0'][23/31] [tp:3/3][dp:5/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s25b1n0'][31/31] [tp:3/3][dp:7/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s19b1n0'][15/31] [tp:3/3][dp:3/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s1b0n0'][17/31] [tp:1/3][dp:4/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s13b0n0'][ 1/31] [tp:1/3][dp:0/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s19b0n0'][ 9/31] [tp:1/3][dp:2/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s25b0n0'][25/31] [tp:1/3][dp:6/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s19b1n0'][13/31] [tp:1/3][dp:3/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s25b1n0'][29/31] [tp:1/3][dp:7/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s1b1n0'][21/31] [tp:1/3][dp:5/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s13b1n0'][ 5/31] [tp:1/3][dp:1/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s19b0n0'][ 8/31] [tp:0/3][dp:2/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s25b0n0'][24/31] [tp:0/3][dp:6/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s1b1n0'][20/31] [tp:0/3][dp:5/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s13b1n0'][ 4/31] [tp:0/3][dp:1/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s25b1n0'][28/31] [tp:0/3][dp:7/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s19b1n0'][12/31] [tp:0/3][dp:3/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s13b0n0'][ 2/31] [tp:2/3][dp:0/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s25b1n0'][30/31] [tp:2/3][dp:7/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s1b0n0'][18/31] [tp:2/3][dp:4/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s13b0n0'][ 0/31] [tp:0/3][dp:0/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s1b0n0'][16/31] [tp:0/3][dp:4/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s19b1n0'][14/31] [tp:2/3][dp:3/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s13b1n0'][ 6/31] [tp:2/3][dp:1/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s1b1n0'][22/31] [tp:2/3][dp:5/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s19b0n0'][10/31] [tp:2/3][dp:2/7]\n    [2025-02-01 13:57:44][I][ezpz/dist:869] ['x3002c0s25b0n0'][26/31] [tp:2/3][dp:6/7]\n    [2025-02-01 13:57:44][I][examples/fsdp_tp:177:__main__] Device mesh created:\n    device_mesh=DeviceMesh([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23], [24, 25, 26, 27], [28, 29, 30, 31]], mesh_dim_names=('dp', 'tp'))\n    [2025-02-01 13:57:44][I][examples/fsdp_tp:187:__main__] config:\n    ModelArgs(dim=4096, n_layers=32, n_heads=32, n_kv_heads=4, vocab_size=32000, multiple_of=360, ffn_dim_multiplier=None, norm_eps=1e-05, max_batch_size=32, max_seq_len=32768, depth_init=True)\n    [2025-02-01 13:58:29][I][examples/fsdp_tp:191:__main__]\n    =================================================================\n    Layer (type:depth-idx)                   Param #\n    =================================================================\n    Transformer                              --\n    ├─Embedding: 1-1                         131,072,000\n    ├─ModuleList: 1-2                        --\n    │    └─TransformerBlock: 2-1             174,891,008\n    │    └─TransformerBlock: 2-2             174,891,008\n    │    └─TransformerBlock: 2-3             174,891,008\n    │    └─TransformerBlock: 2-4             174,891,008\n    │    └─TransformerBlock: 2-5             174,891,008\n    │    └─TransformerBlock: 2-6             174,891,008\n    │    └─TransformerBlock: 2-7             174,891,008\n    │    └─TransformerBlock: 2-8             174,891,008\n    │    └─TransformerBlock: 2-9             174,891,008\n    │    └─TransformerBlock: 2-10            174,891,008\n    │    └─TransformerBlock: 2-11            174,891,008\n    │    └─TransformerBlock: 2-12            174,891,008\n    │    └─TransformerBlock: 2-13            174,891,008\n    │    └─TransformerBlock: 2-14            174,891,008\n    │    └─TransformerBlock: 2-15            174,891,008\n    │    └─TransformerBlock: 2-16            174,891,008\n    │    └─TransformerBlock: 2-17            174,891,008\n    │    └─TransformerBlock: 2-18            174,891,008\n    │    └─TransformerBlock: 2-19            174,891,008\n    │    └─TransformerBlock: 2-20            174,891,008\n    │    └─TransformerBlock: 2-21            174,891,008\n    │    └─TransformerBlock: 2-22            174,891,008\n    │    └─TransformerBlock: 2-23            174,891,008\n    │    └─TransformerBlock: 2-24            174,891,008\n    │    └─TransformerBlock: 2-25            174,891,008\n    │    └─TransformerBlock: 2-26            174,891,008\n    │    └─TransformerBlock: 2-27            174,891,008\n    │    └─TransformerBlock: 2-28            174,891,008\n    │    └─TransformerBlock: 2-29            174,891,008\n    │    └─TransformerBlock: 2-30            174,891,008\n    │    └─TransformerBlock: 2-31            174,891,008\n    │    └─TransformerBlock: 2-32            174,891,008\n    ├─RMSNorm: 1-3                           4,096\n    ├─Linear: 1-4                            131,072,000\n    =================================================================\n    Total params: 5,858,660,352\n    Trainable params: 5,858,660,352\n    Non-trainable params: 0\n    =================================================================\n    [2025-02-01 13:58:33][I][examples/fsdp_tp:163:__main__] Model after parallelization:\n    sharded_model=FullyShardedDataParallel(\n      (_fsdp_wrapped_module): Transformer(\n        (tok_embeddings): Embedding(32000, 4096)\n        (layers): ModuleList(\n          (0-31): 32 x TransformerBlock(\n            (attention): Attention(\n              (wq): Linear(in_features=4096, out_features=4096, bias=False)\n              (wk): Linear(in_features=4096, out_features=512, bias=False)\n              (wv): Linear(in_features=4096, out_features=512, bias=False)\n              (wo): Linear(in_features=4096, out_features=4096, bias=False)\n            )\n            (feed_forward): FeedForward(\n              (w1): Linear(in_features=4096, out_features=11160, bias=False)\n              (w2): Linear(in_features=11160, out_features=4096, bias=False)\n              (w3): Linear(in_features=4096, out_features=11160, bias=False)\n            )\n            (attention_norm): RMSNorm()\n            (ffn_norm): RMSNorm()\n          )\n        )\n        (norm): RMSNorm()\n        (output): Linear(in_features=4096, out_features=32000, bias=False)\n      )\n    )\n    \n    [2025-02-01 13:58:33][I][examples/fsdp_tp:196:__main__] Creating optimizer=AdamW with lr=0.003\n    [2025-02-01 13:58:33][I][examples/fsdp_tp:245:__main__] Starting 2D training...\n    [2025-02-01 13:58:43][I][examples/fsdp_tp:281:__main__] epoch=0 iter=0 loss=10.831984 dt=9.846861 dtf=5.257388 dtb=4.589474\n    [2025-02-01 13:58:43][I][examples/fsdp_tp:294:__main__] inp.shape=torch.Size([4, 2048])\n    [2025-02-01 13:58:50][I][examples/fsdp_tp:281:__main__] epoch=0 iter=1 loss=14.379510 dt=3.133007 dtf=2.497247 dtb=0.635759\n    [2025-02-01 13:58:57][I][examples/fsdp_tp:281:__main__] epoch=0 iter=2 loss=12.807954 dt=3.584346 dtf=2.438942 dtb=1.145403\n     [2025-02-01 13:59:04][I][examples/fsdp_tp:281:__main__] epoch=0 iter=3 loss=14.743041 dt=3.174036 dtf=2.150096 dtb=1.023941\n    [2025-02-01 13:59:12][I][examples/fsdp_tp:281:__main__] epoch=0 iter=4 loss=17.754011 dt=4.312059 dtf=2.639566 dtb=1.672493\n    [2025-02-01 13:59:20][I][examples/fsdp_tp:281:__main__] epoch=0 iter=5 loss=20.464378 dt=3.201004 dtf=2.256773 dtb=0.944231\n    [2025-02-01 13:59:28][I][examples/fsdp_tp:281:__main__] epoch=0 iter=6 loss=24.087290 dt=4.283923 dtf=2.834405 dtb=1.449518\n    [2025-02-01 13:59:35][I][examples/fsdp_tp:281:__main__] epoch=0 iter=7 loss=25.551876 dt=3.148789 dtf=2.362422 dtb=0.786367\n    [2025-02-01 13:59:43][I][examples/fsdp_tp:281:__main__] epoch=0 iter=8 loss=26.223629 dt=4.989230 dtf=3.239145 dtb=1.750085\n    [2025-02-01 13:59:51][I][examples/fsdp_tp:281:__main__] epoch=0 iter=9 loss=24.683023 dt=3.376321 dtf=2.469399 dtb=0.906922\n    [2025-02-01 13:59:59][I][examples/fsdp_tp:281:__main__] epoch=0 iter=10 loss=22.594797 dt=4.572377 dtf=2.899021 dtb=1.673356\n    [2025-02-01 14:00:07][I][examples/fsdp_tp:281:__main__] epoch=0 iter=11 loss=20.575119 dt=3.941751 dtf=2.874172 dtb=1.067579\n    [2025-02-01 14:00:15][I][examples/fsdp_tp:281:__main__] epoch=0 iter=12 loss=20.123554 dt=4.221992 dtf=3.029075 dtb=1.192917\n    [2025-02-01 14:00:23][I][examples/fsdp_tp:281:__main__] epoch=0 iter=13 loss=26.296919 dt=3.415069 dtf=2.524659 dtb=0.890410\n    [2025-02-01 14:00:32][I][examples/fsdp_tp:281:__main__] epoch=0 iter=14 loss=27.970491 dt=4.739544 dtf=2.963098 dtb=1.776446\n    [2025-02-01 14:00:39][I][examples/fsdp_tp:281:__main__] epoch=0 iter=15 loss=20.901236 dt=3.166723 dtf=2.169875 dtb=0.996849\n    [2025-02-01 14:00:47][I][examples/fsdp_tp:281:__main__] epoch=0 iter=16 loss=21.884733 dt=4.412706 dtf=2.337634 dtb=2.075072\n    [2025-02-01 14:00:55][I][examples/fsdp_tp:281:__main__] epoch=0 iter=17 loss=24.030357 dt=3.159311 dtf=2.286639 dtb=0.872672\n    [2025-02-01 14:01:03][I][examples/fsdp_tp:281:__main__] epoch=0 iter=18 loss=27.559183 dt=4.385118 dtf=2.503596 dtb=1.881523\n    [2025-02-01 14:01:11][I][examples/fsdp_tp:281:__main__] epoch=0 iter=19 loss=29.530497 dt=3.034728 dtf=2.087179 dtb=0.947550\n    [2025-02-01 14:01:19][I][examples/fsdp_tp:281:__main__] epoch=0 iter=20 loss=31.549356 dt=4.096217 dtf=2.819700 dtb=1.276518\n    [2025-02-01 14:01:27][I][examples/fsdp_tp:281:__main__] epoch=0 iter=21 loss=34.394360 dt=3.153647 dtf=2.353484 dtb=0.800163\n    [2025-02-01 14:01:35][I][examples/fsdp_tp:281:__main__] epoch=0 iter=22 loss=35.246937 dt=4.501455 dtf=2.857567 dtb=1.643888\n    [2025-02-01 14:01:43][I][examples/fsdp_tp:281:__main__] epoch=0 iter=23 loss=36.298687 dt=3.637440 dtf=2.376551 dtb=1.260889\n    [2025-02-01 14:01:51][I][examples/fsdp_tp:281:__main__] epoch=0 iter=24 loss=39.277588 dt=4.178340 dtf=2.414844 dtb=1.763496\n    [2025-02-01 14:01:59][I][examples/fsdp_tp:281:__main__] epoch=0 iter=25 loss=40.016758 dt=3.294702 dtf=2.437459 dtb=0.857243\n    ```\n    \n    \u003c/details\u003e\n    \n    \n    \u003c/details\u003e\n\n### 🦙 Llama3\n\n- See [`src/mmm/train.py`](/src/mmm/train.py)\n\n1. Download data:\n\n    ```bash\n    python3 src/mmm/data/download_tokenizer.py \\\n        --repo_id meta-llama/Meta-Llama-3.1-8B \\\n        --tokenizer_path \"original\"\n    ```\n\n2. Launch training:\n\n    ```bash\n    launch python3 -m mmm.train --job.config_file train_configs/llama3_8b.toml --training.seq_len=2048\n    ```\n\n    \u003cdetails closed\u003e\u003csummary\u003eOutput:\u003c/summary\u003e\n    \n    \u003cdetails closed\u003e\u003csummary\u003ePolaris @ ALCF:\u003c/summary\u003e\n    \n    ```bash\n    (👻 2024-04-29)\n    #[09:18:59 AM][x3005c0s37b1n0][/e/a/f/p/s/mmm][🌱 main][📝🤷✓] [⏱️ 2m23s]\n    $ launch python3 -Wignore -m mmm.train --job.config_file train_configs/llama3_8b.toml --training.seq_len=2048 | tee train-llama3-8b.log\n    ```\n    \n    ```python\n    Connected to tcp://x3005c0s37b1n0.hsn.cm.polaris.alcf.anl.gov:7919\n    Found executable /eagle/argonne_tpc/foremans/projects/saforem2/mmm/venvs/2024-04-29/bin/python3\n    Launching application a95db38c-5a1f-4f20-a362-6cd48cc0f009\n    Using PMI port 40193,40194\n    [2025-02-03 09:19:25][I][datasets/config:54:datasets] PyTorch version 2.6.0+cu124 available.\n    [2025-02-03 09:19:25][I][datasets/config:112:datasets] TensorFlow version 2.16.1 available.\n    [2025-02-03 09:19:25][I][datasets/config:125:datasets] JAX version 0.4.26 available.\n    [2025-02-03 09:19:41][I][mmm/train:71:__main__] Starting job: Llama 3 8B training\n    [rank4]:[W203 09:19:41.475782740 ProcessGroupNCCL.cpp:4561] [PG ID 0 PG GUID 0 Rank 4]  using GPU 0 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.\n    [rank0]:[W203 09:19:41.501737516 ProcessGroupNCCL.cpp:4561] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.\n    [rank7]:[W203 09:19:42.070718085 ProcessGroupNCCL.cpp:4561] [PG ID 0 PG GUID 0 Rank 7]  using GPU 3 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.\n    [rank5]:[W203 09:19:42.071337647 ProcessGroupNCCL.cpp:4561] [PG ID 0 PG GUID 0 Rank 5]  using GPU 1 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.\n    [rank6]:[W203 09:19:42.071862543 ProcessGroupNCCL.cpp:4561] [PG ID 0 PG GUID 0 Rank 6]  using GPU 2 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.\n    [rank2]:[W203 09:19:42.067891022 ProcessGroupNCCL.cpp:4561] [PG ID 0 PG GUID 0 Rank 2]  using GPU 2 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.\n    [rank3]:[W203 09:19:42.068768126 ProcessGroupNCCL.cpp:4561] [PG ID 0 PG GUID 0 Rank 3]  using GPU 3 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.\n    [rank1]:[W203 09:19:42.069066595 ProcessGroupNCCL.cpp:4561] [PG ID 0 PG GUID 0 Rank 1]  using GPU 1 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in ba\n    rrier() to force use of a particular device, or call init_process_group() with a device_id.\n    [2025-02-03 09:19:43][I][ezpz/dist:823] Using device='cuda' with backend='DDP' + 'nccl' for distributed training.\n    [2025-02-03 09:19:43][I][ezpz/dist:869] ['x3005c0s37b1n0'][2/7]\n    [2025-02-03 09:19:43][I][ezpz/dist:869] ['x3005c0s37b1n0'][1/7]\n    [2025-02-03 09:19:43][I][ezpz/dist:869] ['x3005c0s7b0n0'][4/7]\n    [2025-02-03 09:19:43][I][ezpz/dist:869] ['x3005c0s7b0n0'][6/7]\n    [2025-02-03 09:19:43][I][ezpz/dist:869] ['x3005c0s7b0n0'][5/7]\n    [2025-02-03 09:19:43][I][ezpz/dist:869] ['x3005c0s37b1n0'][3/7]\n    [2025-02-03 09:19:43][I][ezpz/dist:869] ['x3005c0s7b0n0'][7/7]\n    [2025-02-03 09:19:43][I][ezpz/dist:869] ['x3005c0s37b1n0'][0/7]\n    \n    [2025-02-03 09:19:43][I][mmm/utils:298] Process group already initialized, skipping init_process_group\n    [2025-02-03 09:19:43][W][mmm/utils:361] Peak flops undefined for: cuda, fallback to A100\n    [2025-02-03 09:19:43][I][mmm/train:113:__main__] Peak FLOPS used for computing MFU: 3.120e+14\n    [2025-02-03 09:19:43][I][parallelisms/parallel_dims:80:mmm.parallelisms.parallel_dims] Building 1-D device mesh with ['dp_shard'], [8]\n    [2025-02-03 09:19:43][I][tokenizer/__init__:18:mmm.data.tokenizer] Building tiktoken tokenizer locally from ./tests/assets/test_tiktoken.model\n    [2025-02-03 09:19:44][I][tokenizer/tiktoken:92:mmm.data.tokenizer.tiktoken] TikTokenizer built: #words 128256, BOS ID 128000, EOS ID 128001\n    [2025-02-03 09:19:44][I][data/hf_datasets:73:mmm.data.hf_datasets] Preparing c4 dataset from allenai/c4\n    [2025-02-03 09:19:49][I][mmm/train:165:__main__] Building llama3 8B with ModelArgs(dim=4096, n_layers=32, n_heads=32, n_kv_heads=8, vocab_size=128256, multiple_of=1024, ffn_dim_multiplier=1.3, norm_eps=1e-05, rope_theta=500000, max_seq_len=2048, depth_init=True, norm_type='rmsnorm')\n    [2025-02-03 09:19:50][I][mmm/train:181:__main__] Model llama3 8B size: 8,030,261,248 total parameters\n    [2025-02-03 09:19:50][I][parallelisms/parallelize_llama:309:mmm.parallelisms.parallelize_llama] Applied selective activation checkpointing to the model\n    [2025-02-03 09:19:50][I][parallelisms/parallelize_llama:114:mmm.parallelisms.parallelize_llama] Applied FSDP to the model\n    [2025-02-03 09:19:51][E][mmm/train:241:__main__] Error getting memory stats: name 'device_memory_monitor' is not defined\n    [2025-02-03 09:19:51][I][mmm/metrics:128] TensorBoard logging enabled. Logs will be saved at ./outputs/tb/20250203-0919\n    [2025-02-03 09:19:51][E][mmm/train:296:__main__] Error resetting memory stats: name 'device_memory_monitor' is not defined\n    [2025-02-03 09:19:51][I][mmm/train:300:__main__] Training starts at step 1, with local batch size 1, global batch size 8, sequence length 2048, total steps 1000 (warmup 200)\n    [2025-02-03 09:19:51][I][mmm/profiling:59] Profiling active. Traces will be saved at ./outputs/profile_trace\n    [2025-02-03 09:20:14][E][mmm/train:438:__main__] Error getting memory stats: name 'device_memory_monitor' is not defined\n    [2025-02-03 09:20:14][I][mmm/train:458:__main__] step=1 global_avg_loss=12.269001 global_max_loss=12.310555 throughput(tps)=86.019928 mfu(%)=1.330297 end_to_end(s)=2.380844 data_loading(s)=0.900774 data_loading(%)=3.783424\n    [2025-02-03 09:20:14][E][mmm/train:498:__main__] Error resetting memory stats: name 'device_memory_monitor' is not defined\n    [2025-02-03 09:20:14][I][mmm/utils:171] Synchronizing and adjusting timeout for all ProcessGroups to 0:01:40\n    \n    [2025-02-03 09:22:59][I][mmm/train:458:__main__] step=10 global_avg_loss=11.101408 global_max_loss=12.190623 throughput(tps)=112.290594 mfu(%)=1.736573 end_to_end(s)=16.414554 data_loading(s)=0.003078 data_loading(%)=0.016878\n    [2025-02-03 09:26:05][I][mmm/train:458:__main__] step=20 global_avg_loss=9.304403 global_max_loss=10.041307 throughput(tps)=110.054333 mfu(%)=1.701989 end_to_end(s)=18.608990 data_loading(s)=0.002827 data_loading(%)=0.015190\n    [2025-02-03 09:29:13][I][mmm/train:458:__main__] step=30 global_avg_loss=8.363521 global_max_loss=10.094500 throughput(tps)=109.026794 mfu(%)=1.686098 end_to_end(s)=18.784373 data_loading(s)=0.003182 data_loading(%)=0.016937\n    [2025-02-03 09:32:18][I][mmm/train:458:__main__] step=40 global_avg_loss=7.754060 global_max_loss=10.717106 throughput(tps)=110.327183 mfu(%)=1.706209 end_to_end(s)=18.562968 data_loading(s)=0.003432 data_loading(%)=0.018489\n    [2025-02-03 09:35:28][I][mmm/train:458:__main__] step=50 global_avg_loss=7.405908 global_max_loss=11.071878 throughput(tps)=108.140150 mfu(%)=1.672386 end_to_end(s)=18.938387 data_loading(s)=0.002782 data_loading(%)=0.014690\n    [2025-02-03 09:38:32][I][mmm/train:458:__main__] step=60 global_avg_loss=7.379553 global_max_loss=16.768843 throughput(tps)=111.114522 mfu(%)=1.718385 end_to_end(s)=18.431434 data_loading(s)=0.002796 data_loading(%)=0.015172\n    [2025-02-03 09:41:36][I][mmm/train:458:__main__] step=70 global_avg_loss=7.180877 global_max_loss=8.136713 throughput(tps)=111.323097 mfu(%)=1.721611 end_to_end(s)=18.396901 data_loading(s)=0.003316 data_loading(%)=0.018022\n    [2025-02-03 09:44:40][I][mmm/train:458:__main__] step=80 global_avg_loss=6.949995 global_max_loss=8.055350 throughput(tps)=111.244684 mfu(%)=1.720398 end_to_end(s)=18.409868 data_loading(s)=0.003094 data_loading(%)=0.016804\n    [2025-02-03 09:47:47][I][mmm/train:458:__main__] step=90 global_avg_loss=6.984058 global_max_loss=8.244547 throughput(tps)=109.736104 mfu(%)=1.697068 end_to_end(s)=18.662955 data_loading(s)=0.003001 data_loading(%)=0.016082\n    [2025-02-03 09:50:52][I][mmm/train:458:__main__] step=100 global_avg_loss=6.869817 global_max_loss=8.179618 throughput(tps)=110.670108 mfu(%)=1.711512 end_to_end(s)=18.505449 data_loading(s)=0.008446 data_loading(%)=0.045638\n    ```\n    \n    \u003c/details\u003e\n    \n    \u003c/details\u003e\n\n### 🖼️ ViT\n\nSee [`src/mmm/examples/vit.py`](/src/mmm/examples/vit.py):\n\n- Launch:\n\n    ```bash\n    launch python3 -m mmm.examples.vit --max_iters 10\n    ```\n    \n    \u003cdetails closed\u003e\u003csummary\u003eOutput:\u003c/summary\u003e\n    \n    \u003cdetails closed\u003e\u003csummary\u003eAurora @ ALCF:\u003c/summary\u003e\n    \n    ```bash\n    [🐍 aurora_nre_models_frameworks-2024.2.1_u1](👻 aurora_nre_models_frameworks-2024.2.1_u1)\n    #[10:02:16 AM][x4404c5s2b0n0][/f/A/f/p/s/mmm][🌱 main][$!?]\n    $ launch python3 -Wignore -m mmm.examples.vit --max_iters 10\n    ```\n    \n    ```python\n    Disabling local launch: multi-node application\n    Connected to tcp://x4509c0s2b0n0.hostmgmt2509.cm.aurora.alcf.anl.gov:7919\n    Found executable /flare/Aurora_deployment/foremans/projects/saforem2/mmm/venvs/aurora_nre_models_frameworks-2024.2.1_u1/bin/python3\n    Launching application 8ac3f76b-2b75-45c0-925d-3bfa3eee7633\n    [2025-01-03 10:04:41,306] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to xpu (auto detect)\n    # ...[clipped]...\n    [2025-01-03 10:04:42,574] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to xpu (auto detect)\n    [2025-01-03 10:05:03.848363][INFO][vit.py:246] - Using native for SDPA backend\n    [2025-01-03 10:05:03.851500][INFO][vit.py:272] - Using AttentionBlock Attention with args.compile=False\n    [2025-01-03 10:05:03.918476][INFO][__init__.py:146] - \u003e initializing tensor parallel with size 1\n    [2025-01-03 10:05:03.918996][INFO][__init__.py:151] - \u003e initializing context parallel with size 1\n    [2025-01-03 10:05:03.919362][INFO][__init__.py:156] - \u003e initializing pipeline with size 1\n    [2025-01-03 10:05:03.919692][INFO][__init__.py:159] - \u003e initializing ddp with size 48\n    2025:01:03-10:05:03:(94856) |CCL_WARN| value of CCL_LOG_LEVEL changed to be error (default:warn)\n    [2025-01-03 10:05:05.399317][INFO][dist.py:818] - Using device='xpu' with backend='DDP' + 'ccl' for distributed training.\n    [2025-01-03 10:05:05.400019][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 0/47]\n    [2025-01-03 10:05:05.399397][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 1/47]\n    [2025-01-03 10:05:05.399353][INFO][dist.py:859] - ['x4509c0s2b0n0'][11/47]\n    [2025-01-03 10:05:05.399368][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 2/47]\n    [2025-01-03 10:05:05.399384][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 3/47]\n    [2025-01-03 10:05:05.399373][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 4/47]\n    [2025-01-03 10:05:05.399388][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 5/47]\n    [2025-01-03 10:05:05.399352][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 7/47]\n    [2025-01-03 10:05:05.399348][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 8/47]\n    [2025-01-03 10:05:05.399372][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 9/47]\n    [2025-01-03 10:05:05.399356][INFO][dist.py:859] - ['x4509c0s2b0n0'][10/47]\n    [2025-01-03 10:05:05.399312][INFO][dist.py:859] - ['x4404c5s2b0n0'][36/47]\n    [2025-01-03 10:05:05.399312][INFO][dist.py:859] - ['x4206c2s1b0n0'][24/47]\n    [2025-01-03 10:05:05.399333][INFO][dist.py:859] - ['x4206c0s0b0n0'][20/47]\n    [2025-01-03 10:05:05.399358][INFO][dist.py:859] - ['x4206c0s0b0n0'][12/47]\n    [2025-01-03 10:05:05.399358][INFO][dist.py:859] - ['x4206c0s0b0n0'][13/47]\n    [2025-01-03 10:05:05.399352][INFO][dist.py:859] - ['x4206c0s0b0n0'][14/47]\n    [2025-01-03 10:05:05.399325][INFO][dist.py:859] - ['x4206c2s1b0n0'][25/47]\n    [2025-01-03 10:05:05.399330][INFO][dist.py:859] - ['x4404c5s2b0n0'][37/47]\n    [2025-01-03 10:05:05.399363][INFO][dist.py:859] - ['x4206c0s0b0n0'][15/47]\n    [2025-01-03 10:05:05.399311][INFO][dist.py:859] - ['x4206c2s1b0n0'][26/47]\n    [2025-01-03 10:05:05.399333][INFO][dist.py:859] - ['x4404c5s2b0n0'][39/47]\n    [2025-01-03 10:05:05.399363][INFO][dist.py:859] - ['x4206c0s0b0n0'][16/47]\n    [2025-01-03 10:05:05.399333][INFO][dist.py:859] - ['x4206c2s1b0n0'][27/47]\n    [2025-01-03 10:05:05.399332][INFO][dist.py:859] - ['x4404c5s2b0n0'][41/47]\n    [2025-01-03 10:05:05.399386][INFO][dist.py:859] - ['x4206c0s0b0n0'][17/47]\n    [2025-01-03 10:05:05.399302][INFO][dist.py:859] - ['x4206c2s1b0n0'][28/47]\n    [2025-01-03 10:05:05.399306][INFO][dist.py:859] - ['x4404c5s2b0n0'][43/47]\n    [2025-01-03 10:05:05.399348][INFO][dist.py:859] - ['x4206c0s0b0n0'][19/47]\n    [2025-01-03 10:05:05.399351][INFO][dist.py:859] - ['x4206c0s0b0n0'][21/47]\n    [2025-01-03 10:05:05.399333][INFO][dist.py:859] - ['x4206c0s0b0n0'][22/47]\n    [2025-01-03 10:05:05.399444][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 6/47]\n    [2025-01-03 10:05:05.399355][INFO][dist.py:859] - ['x4206c0s0b0n0'][23/47]\n    [2025-01-03 10:05:05.399294][INFO][dist.py:859] - ['x4404c5s2b0n0'][44/47]\n    [2025-01-03 10:05:05.399317][INFO][dist.py:859] - ['x4206c2s1b0n0'][29/47]\n    [2025-01-03 10:05:05.399405][INFO][dist.py:859] - ['x4206c0s0b0n0'][18/47]\n    [2025-01-03 10:05:05.399306][INFO][dist.py:859] - ['x4404c5s2b0n0'][45/47]\n    [2025-01-03 10:05:05.399297][INFO][dist.py:859] - ['x4206c2s1b0n0'][31/47]\n    [2025-01-03 10:05:05.399296][INFO][dist.py:859] - ['x4404c5s2b0n0'][46/47]\n    [2025-01-03 10:05:05.399286][INFO][dist.py:859] - ['x4206c2s1b0n0'][32/47]\n    [2025-01-03 10:05:05.399320][INFO][dist.py:859] - ['x4206c2s1b0n0'][33/47]\n    [2025-01-03 10:05:05.399289][INFO][dist.py:859] - ['x4206c2s1b0n0'][34/47]\n    [2025-01-03 10:05:05.399290][INFO][dist.py:859] - ['x4404c5s2b0n0'][47/47]\n    [2025-01-03 10:05:05.399290][INFO][dist.py:859] - ['x4206c2s1b0n0'][35/47]\n    [2025-01-03 10:05:05.399323][INFO][dist.py:859] - ['x4404c5s2b0n0'][38/47]\n    [2025-01-03 10:05:05.399324][INFO][dist.py:859] - ['x4206c2s1b0n0'][30/47]\n    [2025-01-03 10:05:05.399324][INFO][dist.py:859] - ['x4404c5s2b0n0'][40/47]\n    [2025-01-03 10:05:05.399373][INFO][dist.py:859] - ['x4404c5s2b0n0'][42/47]\n    [2025-01-03 10:05:05.404629][INFO][vit.py:83] - config=ViTConfig(img_size=224, batch_size=128, num_heads=16, head_dim=64, depth=24, patch_size=16)\n    [2025-01-03 10:05:33.726284][INFO][vit.py:122] -\n    ==========================================================================================\n    Layer (type:depth-idx)                   Output Shape              Param #\n    ==========================================================================================\n    VisionTransformer                        [128, 1000]               200,704\n    ├─PatchEmbed: 1-1                        [128, 196, 1024]          787,456\n    ├─Dropout: 1-2                           [128, 196, 1024]          --\n    ├─Identity: 1-3                          [128, 196, 1024]          --\n    ├─Identity: 1-4                          [128, 196, 1024]          --\n    ├─Sequential: 1-5                        [128, 196, 1024]          302,235,648\n    ├─Identity: 1-6                          [128, 196, 1024]          --\n    ├─LayerNorm: 1-7                         [128, 1024]               2,048\n    ├─Dropout: 1-8                           [128, 1024]               --\n    ├─Linear: 1-9                            [128, 1000]               1,025,000\n    ==========================================================================================\n    Total params: 304,250,856\n    Trainable params: 304,250,856\n    Non-trainable params: 0\n    Total mult-adds (G): 58.57\n    ==========================================================================================\n    Input size (MB): 77.07\n    Forward/backward pass size (MB): 54465.11\n    Params size (MB): 1216.20\n    Estimated Total Size (MB): 55758.38\n    ==========================================================================================\n    [2025-01-03 10:05:33.838873][INFO][vit.py:138] - Model size: nparams=1.42 B\n    [2025-01-03 10:05:33.875672][INFO][vit.py:163] - Training with 48 x xpu (s), using torch_dtype=torch.bfloat16\n    [2025-01-03 10:05:51.652331][INFO][vit.py:181] - iter=0 loss=7.110761 dt=17.446471 dtf=16.439806 dtb=0.984764\n    [2025-01-03 10:05:52.375636][INFO][vit.py:181] - iter=1 loss=7.161612 dt=0.699467 dtf=0.047248 dtb=0.641429\n    [2025-01-03 10:05:53.105314][INFO][vit.py:181] - iter=2 loss=7.230668 dt=0.704598 dtf=0.048066 dtb=0.645473\n    [2025-01-03 10:05:53.825272][INFO][vit.py:181] - iter=3 loss=7.113066 dt=0.695661 dtf=0.042115 dtb=0.642792\n    [2025-01-03 10:05:54.546189][INFO][vit.py:181] - iter=4 loss=6.960967 dt=0.697206 dtf=0.042410 dtb=0.643944\n    [2025-01-03 10:05:55.266622][INFO][vit.py:181] - iter=5 loss=6.963243 dt=0.696353 dtf=0.044059 dtb=0.641429\n    [2025-01-03 10:05:55.988250][INFO][vit.py:181] - iter=6 loss=7.016630 dt=0.696779 dtf=0.043519 dtb=0.642242\n    [2025-01-03 10:05:56.706743][INFO][vit.py:181] - iter=7 loss=7.080251 dt=0.693541 dtf=0.043349 dtb=0.639498\n    [2025-01-03 10:05:57.426416][INFO][vit.py:181] - iter=8 loss=7.020833 dt=0.695024 dtf=0.043666 dtb=0.640416\n    [2025-01-03 10:05:58.147060][INFO][vit.py:181] - iter=9 loss=7.047978 dt=0.696053 dtf=0.045267 dtb=0.639672\n    [2025-01-03 10:05:58.865314][INFO][vit.py:181] - iter=10 loss=7.021929 dt=0.693284 dtf=0.043635 dtb=0.638591\n    [2025-01-03 10:05:59.218314][INFO][history.py:723] - Saving train_iter plot to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/plots/mplot\n    [2025-01-03 10:05:59.499058][INFO][history.py:723] - Saving train_loss plot to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/plots/mplot\n    [2025-01-03 10:05:59.740021][INFO][history.py:723] - Saving train_dt plot to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/plots/mplot\n    [2025-01-03 10:05:59.984435][INFO][history.py:723] - Saving train_dtf plot to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/plots/mplot\n    [2025-01-03 10:06:00.214398][INFO][history.py:723] - Saving train_dtb plot to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/plots/mplot\n    [2025-01-03 10:06:00.433872][INFO][history.py:603] - Saving tplots to /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/plots/tplot\n                            train_iter [2025-01-03-100600]\n        ┌─────────────────────────────────────────────────────────────────────┐\n    10.0┤                                                                  ▗▄▞│\n        │                                                             ▗▄▄▀▀▘  │\n        │                                                          ▗▄▀▘       │\n     8.3┤                                                       ▄▄▀▘          │\n        │                                                  ▗▄▄▀▀              │\n        │                                              ▗▄▀▀▘                  │\n     6.7┤                                           ▄▞▀▘                      │\n        │                                       ▄▄▀▀                          │\n     5.0┤                                  ▗▄▄▀▀                              │\n        │                               ▄▄▀▘                                  │\n        │                           ▗▄▞▀                                      │\n     3.3┤                       ▄▄▞▀▘                                         │\n        │                   ▄▞▀▀                                              │\n        │               ▗▄▀▀                                                  │\n     1.7┤           ▗▄▞▀▘                                                     │\n        │       ▄▄▞▀▘                                                         │\n        │   ▗▄▞▀                                                              │\n     0.0┤▄▄▀▘                                                                 │\n        └┬──────┬──────┬─────┬──────┬──────┬──────┬──────┬─────┬──────┬───────┘\n         1      2      3     4      5      6      7      8     9     10\n    train_iter\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/plots/tplot/train_iter.txt\n                             train_loss [2025-01-03-100600]\n         ┌────────────────────────────────────────────────────────────────────┐\n    7.231┤             ▟                                                      │\n         │           ▗▞ ▚                                                     │\n         │          ▄▘   ▚                                                    │\n    7.186┤        ▗▞      ▚                                                   │\n         │       ▄▘        ▚                                                  │\n         │     ▄▀           ▚                                                 │\n    7.141┤   ▄▀              ▚                                                │\n         │ ▄▀                 ▚                                               │\n    7.096┤▀                   ▝▖                                              │\n         │                     ▚                                              │\n         │                      ▌                        ▞▖                   │\n    7.051┤                      ▝▖                     ▄▀ ▝▚▖                 │\n         │                       ▚                   ▗▀     ▝▚▖      ▄▄▚▄▖    │\n         │                        ▌                ▗▞▘        ▝▚▄▄▞▀▀    ▝▀▚▄▄│\n    7.006┤                        ▝▖              ▞▘                          │\n         │                         ▚            ▄▀                            │\n         │                          ▌         ▄▀                              │\n    6.961┤                          ▝▄▄▄▄▄▄▄▄▀                                │\n         └┬──────┬─────┬──────┬──────┬──────┬─────┬──────┬──────┬─────┬───────┘\n          1      2     3      4      5      6     7      8      9    10\n    train_loss\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/plots/tplot/train_loss.txt\n                             train_dt [2025-01-03-100600]\n        ┌─────────────────────────────────────────────────────────────────────┐\n    17.4┤▌                                                                    │\n        │▚                                                                    │\n        │▝▖                                                                   │\n    14.7┤ ▌                                                                   │\n        │ ▐                                                                   │\n        │  ▌                                                                  │\n    11.9┤  ▚                                                                  │\n        │  ▝▖                                                                 │\n     9.1┤   ▌                                                                 │\n        │   ▐                                                                 │\n        │    ▌                                                                │\n     6.3┤    ▚                                                                │\n        │    ▝▖                                                               │\n        │     ▌                                                               │\n     3.5┤     ▐                                                               │\n        │      ▌                                                              │\n        │      ▚                                                              │\n     0.7┤      ▝▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄│\n        └┬──────┬──────┬─────┬──────┬──────┬──────┬──────┬─────┬──────┬───────┘\n         1      2      3     4      5      6      7      8     9     10\n    train_dt\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/plots/tplot/train_dt.txt\n                             train_dtf [2025-01-03-100600]\n        ┌─────────────────────────────────────────────────────────────────────┐\n    16.4┤▌                                                                    │\n        │▚                                                                    │\n        │▝▖                                                                   │\n    13.7┤ ▌                                                                   │\n        │ ▐                                                                   │\n        │  ▌                                                                  │\n    11.0┤  ▚                                                                  │\n        │  ▝▖                                                                 │\n     8.2┤   ▌                                                                 │\n        │   ▐                                                                 │\n        │    ▌                                                                │\n     5.5┤    ▚                                                                │\n        │    ▝▖                                                               │\n        │     ▌                                                               │\n     2.8┤     ▐                                                               │\n        │      ▌                                                              │\n        │      ▚                                                              │\n     0.0┤      ▝▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄│\n        └┬──────┬──────┬─────┬──────┬──────┬──────┬──────┬─────┬──────┬───────┘\n         1      2      3     4      5      6      7      8     9     10\n    train_dtf\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/plots/tplot/train_dtf.txt\n                              train_dtb [2025-01-03-100600]\n         ┌────────────────────────────────────────────────────────────────────┐\n    0.985┤▌                                                                   │\n         │▚                                                                   │\n         │▝▖                                                                  │\n    0.927┤ ▌                                                                  │\n         │ ▐                                                                  │\n         │  ▌                                                                 │\n    0.869┤  ▚                                                                 │\n         │  ▝▖                                                                │\n    0.812┤   ▌                                                                │\n         │   ▐                                                                │\n         │    ▌                                                               │\n    0.754┤    ▚                                                               │\n         │    ▝▖                                                              │\n         │     ▌                                                              │\n    0.696┤     ▐                                                              │\n         │      ▌                                                             │\n         │      ▚                                                             │\n    0.639┤      ▝▄▄▄▄▄▄▞▄▄▄▄▄▄▄▄▄▄▄▄▄▚▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄│\n         └┬──────┬─────┬──────┬──────┬──────┬─────┬──────┬──────┬─────┬───────┘\n          1      2     3      4      5      6     7      8      9    10\n    train_dtb\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/plots/tplot/train_dtb.txt\n    [2025-01-03 10:06:00.618027][INFO][utils.py:132] - Saving dataset to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-vit/train_dataset.h5\n    [2025-01-03 10:06:00.656925][INFO][vit.py:200] - dataset=\u003cxarray.Dataset\u003e Size: 528B\n    Dimensions:     (draw: 11)\n    Coordinates:\n      * draw        (draw) int64 88B 0 1 2 3 4 5 6 7 8 9 10\n    Data variables:\n        train_iter  (draw) int64 88B 0 1 2 3 4 5 6 7 8 9 10\n        train_loss  (draw) float64 88B 7.111 7.162 7.231 7.113 ... 7.021 7.048 7.022\n        train_dt    (draw) float64 88B 17.45 0.6995 0.7046 ... 0.695 0.6961 0.6933\n        train_dtf   (draw) float64 88B 16.44 0.04725 0.04807 ... 0.04527 0.04364\n        train_dtb   (draw) float64 88B 0.9848 0.6414 0.6455 ... 0.6404 0.6397 0.6386\n    [2025-01-03 10:06:00.664331][INFO][vit.py:279] - Took 56.82 seconds\n    Application 8ac3f76b resources: utime=10756s stime=1931s maxrss=4638452KB inblock=1297516 oublock=1920 minflt=38089683 majflt=12671790 nvcsw=1276384 nivcsw=1603670\n    took: 0h:01m:30s\n    ```\n    \n    \u003c/details\u003e\n    \n    \u003cdetails closed\u003e\u003csummary\u003ePolaris @ ALCF:\u003c/summary\u003e\n    \n    Command:\n    \n    ```bash\n    launch python3 -m mmm.examples.vit --max_iters 10\n    ```\n    \n    Output:\n    \n    ```python\n    #[🐍 aurora_nre_models_frameworks-2024.2.1_u1](👻 aurora_nre_models_frameworks-2024.2.1_u1)\n    #[07:48:56 PM][x4603c7s2b0n0][/f/A/f/p/s/mmm][🌱 main][$!?][⏱️ 29s]\n    Disabling local launch: multi-node application\n    Connected to tcp://x4603c7s2b0n0.hostmgmt2603.cm.aurora.alcf.anl.gov:7919\n    Found executable /flare/Aurora_deployment/foremans/projects/saforem2/mmm/venvs/aurora_nre_models_frameworks-2024.2.1_u1/bin/python3\n    Launching application 66c85dca-4341-4edd-8100-9b8181ff3759\n    [2024-12-31 19:51:22.275852][INFO][__init__.py:146] - \u003e initializing tensor parallel with size 1\n    [2024-12-31 19:51:22.278271][INFO][__init__.py:151] - \u003e initializing context parallel with size 1\n    [2024-12-31 19:51:22.278651][INFO][__init__.py:156] - \u003e initializing pipeline with size 1\n    [2024-12-31 19:51:22.278991][INFO][__init__.py:159] - \u003e initializing ddp with size 24\n    2024:12:31-19:51:22:(148734) |CCL_WARN| value of CCL_LOG_LEVEL changed to be error (default:warn)\n    [2024-12-31 19:51:23.535446][INFO][dist.py:810] - Using device='xpu' with backend='DDP' + 'ccl' for distributed training.\n    [2024-12-31 19:51:23.536194][INFO][dist.py:848] - ['x4603c7s2b0n0'][ 0/23]\n    [2024-12-31 19:51:23.535458][INFO][dist.py:848] - ['x4603c7s2b0n0'][ 6/23]\n    [2024-12-31 19:51:23.535450][INFO][dist.py:848] - ['x4603c7s2b0n0'][ 8/23]\n    [2024-12-31 19:51:23.535448][INFO][dist.py:848] - ['x4603c7s2b0n0'][11/23]\n    [2024-12-31 19:51:23.535473][INFO][dist.py:848] - ['x4603c7s2b0n0'][ 1/23]\n    [2024-12-31 19:51:23.535474][INFO][dist.py:848] - ['x4603c7s2b0n0'][ 2/23]\n    [2024-12-31 19:51:23.535490][INFO][dist.py:848] - ['x4603c7s2b0n0'][ 3/23]\n    [2024-12-31 19:51:23.535487][INFO][dist.py:848] - ['x4603c7s2b0n0'][ 4/23]\n    [2024-12-31 19:51:23.535477][INFO][dist.py:848] - ['x4603c7s2b0n0'][ 5/23]\n    [2024-12-31 19:51:23.535449][INFO][dist.py:848] - ['x4603c7s2b0n0'][ 7/23]\n    [2024-12-31 19:51:23.535456][INFO][dist.py:848] - ['x4603c7s2b0n0'][ 9/23]\n    [2024-12-31 19:51:23.535460][INFO][dist.py:848] - ['x4603c7s2b0n0'][10/23]\n    [2024-12-31 19:51:23.535470][INFO][dist.py:848] - ['x4411c6s2b0n0'][12/23]\n    [2024-12-31 19:51:23.535457][INFO][dist.py:848] - ['x4411c6s2b0n0'][13/23]\n    [2024-12-31 19:51:23.535459][INFO][dist.py:848] - ['x4411c6s2b0n0'][14/23]\n    [2024-12-31 19:51:23.535461][INFO][dist.py:848] - ['x4411c6s2b0n0'][15/23]\n    [2024-12-31 19:51:23.535490][INFO][dist.py:848] - ['x4411c6s2b0n0'][16/23]\n    [2024-12-31 19:51:23.535474][INFO][dist.py:848] - ['x4411c6s2b0n0'][17/23]\n    [2024-12-31 19:51:23.535453][INFO][dist.py:848] - ['x4411c6s2b0n0'][19/23]\n    [2024-12-31 19:51:23.535456][INFO][dist.py:848] - ['x4411c6s2b0n0'][20/23]\n    [2024-12-31 19:51:23.535456][INFO][dist.py:848] - ['x4411c6s2b0n0'][21/23]\n    [2024-12-31 19:51:23.535458][INFO][dist.py:848] - ['x4411c6s2b0n0'][22/23]\n    [2024-12-31 19:51:23.535451][INFO][dist.py:848] - ['x4411c6s2b0n0'][23/23]\n    [2024-12-31 19:51:23.535465][INFO][dist.py:848] - ['x4411c6s2b0n0'][18/23]\n    [2024-12-31 19:51:23.781740][INFO][vit.py:250] - Using native for SDPA backend\n    [2024-12-31 19:51:23.782234][INFO][vit.py:276] - Using AttentionBlock Attention with args.compile=False\n    [2024-12-31 19:51:23.782641][INFO][vit.py:86] - config=ViTConfig(img_size=224, batch_size=128, num_heads=16, head_dim=64, depth=24, patch_size=16)\n    [2024-12-31 19:51:53.240524][INFO][vit.py:122] -\n    ==========================================================================================\n    Layer (type:depth-idx)                   Output Shape              Param #\n    ==========================================================================================\n    VisionTransformer                        [128, 1000]               200,704\n    ├─PatchEmbed: 1-1                        [128, 196, 1024]          787,456\n    ├─Dropout: 1-2                           [128, 196, 1024]          --\n    ├─Identity: 1-3                          [128, 196, 1024]          --\n    ├─Identity: 1-4                          [128, 196, 1024]          --\n    ├─Sequential: 1-5                        [128, 196, 1024]          302,235,648\n    ├─Identity: 1-6                          [128, 196, 1024]          --\n    ├─LayerNorm: 1-7                         [128, 1024]               2,048\n    ├─Dropout: 1-8                           [128, 1024]               --\n    ├─Linear: 1-9                            [128, 1000]               1,025,000\n    ==========================================================================================\n    Total params: 304,250,856\n    Trainable params: 304,250,856\n    Non-trainable params: 0\n    Total mult-adds (G): 58.57\n    ==========================================================================================\n    Input size (MB): 77.07\n    Forward/backward pass size (MB): 54465.11\n    Params size (MB): 1216.20\n    Estimated Total Size (MB): 55758.38\n    ==========================================================================================\n    [2024-12-31 19:51:53.434376][INFO][vit.py:177] - Training with 24 x xpu (s), using torch_dtype=torch.bfloat16\n    [2024-12-31 19:52:09.714576][INFO][vit.py:204] - iter=0 loss=7.121100 dt=15.842163 dtf=14.856410 dtb=0.966353\n    [2024-12-31 19:52:10.656914][INFO][vit.py:204] - iter=1 loss=7.239427 dt=0.918050 dtf=0.299365 dtb=0.606726\n    [2024-12-31 19:52:11.593232][INFO][vit.py:204] - iter=2 loss=7.290485 dt=0.912320 dtf=0.295520 dtb=0.604882\n    [2024-12-31 19:52:12.525405][INFO][vit.py:204] - iter=3 loss=7.131482 dt=0.908596 dtf=0.294570 dtb=0.602297\n    [2024-12-31 19:52:13.454671][INFO][vit.py:204] - iter=4 loss=7.143745 dt=0.905077 dtf=0.290769 dtb=0.602704\n    [2024-12-31 19:52:14.376427][INFO][vit.py:204] - iter=5 loss=7.047148 dt=0.897958 dtf=0.291084 dtb=0.596143\n    [2024-12-31 19:52:15.299474][INFO][vit.py:204] - iter=6 loss=7.120049 dt=0.899810 dtf=0.292079 dtb=0.597195\n    [2024-12-31 19:52:16.218167][INFO][vit.py:204] - iter=7 loss=6.963861 dt=0.895402 dtf=0.293643 dtb=0.591206\n    [2024-12-31 19:52:17.138623][INFO][vit.py:204] - iter=8 loss=7.105645 dt=0.896737 dtf=0.290377 dtb=0.594644\n    [2024-12-31 19:52:18.065521][INFO][vit.py:204] - iter=9 loss=7.102675 dt=0.903153 dtf=0.292257 dtb=0.600275\n    [2024-12-31 19:52:18.992375][INFO][vit.py:204] - iter=10 loss=7.045384 dt=0.902353 dtf=0.291056 dtb=0.600508\n                            train_iter [2024-12-31-195221]\n        ┌─────────────────────────────────────────────────────────────────────┐\n    10.0┤                                                                  ▗▄▞│\n        │                                                             ▗▄▄▀▀▘  │\n        │                                                          ▗▄▀▘       │\n     8.3┤                                                       ▄▄▀▘          │\n        │                                                  ▗▄▄▀▀              │\n        │                                              ▗▄▀▀▘                  │\n     6.7┤                                           ▄▞▀▘                      │\n        │                                       ▄▄▀▀                          │\n     5.0┤                                  ▗▄▄▀▀                              │\n        │                               ▄▄▀▘                                  │\n        │                           ▗▄▞▀                                      │\n     3.3┤                       ▄▄▞▀▘                                         │\n        │                   ▄▞▀▀                                              │\n        │               ▗▄▀▀                                                  │\n     1.7┤           ▗▄▞▀▘                                                     │\n        │       ▄▄▞▀▘                                                         │\n        │   ▗▄▞▀                                                              │\n     0.0┤▄▄▀▘                                                                 │\n        └┬──────┬──────┬─────┬──────┬──────┬──────┬──────┬─────┬──────┬───────┘\n         1      2      3     4      5      6      7      8     9     10\n    train_iter                        train/iter\n    [2024-12-31 19:52:21.800856][INFO][plot.py:220] - Appending plot to: /flare/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/plots/vit/tplot/train_iter.txt\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/plots/vit/tplot/train_iter.txt\n                             train_loss [2024-12-31-195221]\n         ┌────────────────────────────────────────────────────────────────────┐\n    7.290┤            ▗▟                                                      │\n         │          ▄▀▘ ▚                                                     │\n         │       ▄▞▀     ▚                                                    │\n    7.236┤      ▞        ▝▖                                                   │\n         │     ▞          ▝▖                                                  │\n         │    ▞            ▝▖                                                 │\n    7.182┤   ▞              ▚                                                 │\n         │  ▞                ▚                                                │\n    7.127┤ ▞                  ▚▄▄▄▄▄▄▚                                        │\n         │▀                           ▚▖          ▟                           │\n         │                             ▝▄       ▗▀ ▚           ▗▀▀▀▀▀▀▀▚▖     │\n    7.073┤                               ▚    ▗▞▘   ▚         ▗▘        ▝▚▖   │\n         │                                ▀▖ ▄▘     ▝▖       ▗▘           ▝▚▖ │\n         │                                 ▝▀        ▝▖     ▗▘              ▝▀│\n    7.018┤                                            ▝▖    ▞                 │\n         │                                             ▚   ▞                  │\n         │                                              ▚ ▞                   │\n    6.964┤                                               ▜                    │\n         └┬──────┬─────┬──────┬──────┬──────┬─────┬──────┬──────┬─────┬───────┘\n          1      2     3      4      5      6     7      8      9    10\n    train_loss                         train/iter\n    [2024-12-31 19:52:21.834715][INFO][plot.py:220] - Appending plot to: /flare/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/plots/vit/tplot/train_loss.txt\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/plots/vit/tplot/train_loss.txt\n                             train_dt [2024-12-31-195221]\n        ┌─────────────────────────────────────────────────────────────────────┐\n    15.8┤▌                                                                    │\n        │▚                                                                    │\n        │▝▖                                                                   │\n    13.4┤ ▌                                                                   │\n        │ ▐                                                                   │\n        │  ▌                                                                  │\n    10.9┤  ▚                                                                  │\n        │  ▝▖                                                                 │\n     8.4┤   ▌                                                                 │\n        │   ▐                                                                 │\n        │    ▌                                                                │\n     5.9┤    ▚                                                                │\n        │    ▝▖                                                               │\n        │     ▌                                                               │\n     3.4┤     ▐                                                               │\n        │      ▌                                                              │\n        │      ▚                                                              │\n     0.9┤      ▝▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄│\n        └┬──────┬──────┬─────┬──────┬──────┬──────┬──────┬─────┬──────┬───────┘\n         1      2      3     4      5      6      7      8     9     10\n    train_dt                          train/iter\n    [2024-12-31 19:52:21.872382][INFO][plot.py:220] - Appending plot to: /flare/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/plots/vit/tplot/train_dt.txt\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/plots/vit/tplot/train_dt.txt\n                             train_dtf [2024-12-31-195221]\n        ┌─────────────────────────────────────────────────────────────────────┐\n    14.9┤▌                                                                    │\n        │▚                                                                    │\n        │▝▖                                                                   │\n    12.4┤ ▌                                                                   │\n        │ ▐                                                                   │\n        │  ▌                                                                  │\n    10.0┤  ▚                                                                  │\n        │  ▝▖                                                                 │\n     7.6┤   ▌                                                                 │\n        │   ▐                                                                 │\n        │    ▌                                                                │\n     5.1┤    ▚                                                                │\n        │    ▝▖                                                               │\n        │     ▌                                                               │\n     2.7┤     ▐                                                               │\n        │      ▌                                                              │\n        │      ▚                                                              │\n     0.3┤      ▝▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄│\n        └┬──────┬──────┬─────┬──────┬──────┬──────┬──────┬─────┬──────┬───────┘\n         1      2      3     4      5      6      7      8     9     10\n    train_dtf                         train/iter\n    [2024-12-31 19:52:21.896801][INFO][plot.py:220] - Appending plot to: /flare/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/plots/vit/tplot/train_dtf.txt\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/plots/vit/tplot/train_dtf.txt\n                              train_dtb [2024-12-31-195221]\n         ┌────────────────────────────────────────────────────────────────────┐\n    0.966┤▌                                                                   │\n         │▚                                                                   │\n         │▝▖                                                                  │\n    0.904┤ ▌                                                                  │\n         │ ▐                                                                  │\n         │  ▌                                                                 │\n    0.841┤  ▚                                                                 │\n         │  ▝▖                                                                │\n    0.779┤   ▚                                                                │\n         │   ▐                                                                │\n         │    ▌                                                               │\n    0.716┤    ▐                                                               │\n         │    ▝▖                                                              │\n         │     ▚                                                              │\n    0.654┤     ▐                                                              │\n         │      ▌                                                             │\n         │      ▐                                                             │\n    0.591┤       ▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▚▄▄▄▄▄▄▄▄▄▄▄▄▞▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▀▀▀▀▀▀▀│\n         └┬──────┬─────┬──────┬──────┬──────┬─────┬──────┬──────┬─────┬───────┘\n          1      2     3      4      5      6     7      8      9    10\n    train_dtb                          train/iter\n    [2024-12-31 19:52:21.921380][INFO][plot.py:220] - Appending plot to: /flare/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/plots/vit/tplot/train_dtb.txt\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/plots/vit/tplot/train_dtb.txt\n    [2024-12-31 19:52:21.956158][INFO][vit.py:224] - dataset=\u003cxarray.Dataset\u003e Size: 308B\n    Dimensions:     (draw: 11)\n    Coordinates:\n      * draw        (draw) int64 88B 0 1 2 3 4 5 6 7 8 9 10\n    Data variables:\n        train_iter  (draw) float32 44B 0.0 1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0 9.0 10.0\n        train_loss  (draw) float32 44B 7.121 7.239 7.29 7.131 ... 7.106 7.103 7.045\n        train_dt    (draw) float32 44B 15.84 0.9181 0.9123 ... 0.8967 0.9032 0.9024\n        train_dtf   (draw) float32 44B 14.86 0.2994 0.2955 ... 0.2904 0.2923 0.2911\n        train_dtb   (draw) float32 44B 0.9664 0.6067 0.6049 ... 0.5946 0.6003 0.6005\n    [2024-12-31 19:52:21.963271][INFO][vit.py:283] - Took 58.18 seconds\n    Application 66c85dca resources: utime=5443s stime=912s maxrss=4527808KB inblock=755188 oublock=712 minflt=18046931 majflt=6336043 nvcsw=562230 nivcsw=788543\n    took: 0h:01m:15s\n    ```\n    \n    \u003c/details\u003e\n    \n    \u003c/details\u003e\n\n### 📝 FSDP\n\n- See [`src/mmm/examples/fsdp.py`](src/mmm/examples/fsdp.py)\n\n- Launch:\n\n    ```bash\n    launch python3 -m mmm.examples.fsdp\n    ```\n    \n    \u003cdetails closed\u003e\u003csummary\u003eOutput:\u003c/summary\u003e\n    \n    \u003cdetails closed\u003e\u003csummary\u003eAurora @ ALCF:\u003c/summary\u003e\n    \n    Command:\n    \n    ```bash\n    launch python3 -Wignore -m mmm.examples.fsdp\n    ```\n    \n    Output:\n    \n    ```python\n    #[🐍 aurora_nre_models_frameworks-2024.2.1_u1](👻 aurora_nre_models_frameworks-2024.2.1_u1)\n    #[10:15:01 AM][x4404c5s2b0n0][/f/A/f/p/s/mmm][🌱 main][$!?][⏱️ 1m29s]\n    # $ launch python3 -Wignore -m mmm.examples.fsdp --epochs 10\n    Disabling local launch: multi-node application\n    Connected to tcp://x4509c0s2b0n0.hostmgmt2509.cm.aurora.alcf.anl.gov:7919\n    Found executable /flare/Aurora_deployment/foremans/projects/saforem2/mmm/venvs/aurora_nre_models_frameworks-2024.2.1_u1/bin/python3\n    Launching application d4145fbf-be87-460e-ac9a-366b2224df60\n    [2025-01-03 10:16:09,820] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to xpu (auto detect)\n    # ...clipped...\n    [2025-01-03 10:16:11,118] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to xpu (auto detect)\n    [2025-01-03 10:16:16.396841][INFO][__init__.py:146] - \u003e initializing tensor parallel with size 1\n    [2025-01-03 10:16:16.399341][INFO][__init__.py:151] - \u003e initializing context parallel with size 1\n    [2025-01-03 10:16:16.399726][INFO][__init__.py:156] - \u003e initializing pipeline with size 1\n    [2025-01-03 10:16:16.400065][INFO][__init__.py:159] - \u003e initializing ddp with size 48\n    2025:01:03-10:16:16:(99961) |CCL_WARN| value of CCL_LOG_LEVEL changed to be error (default:warn)\n    [2025-01-03 10:16:17.795345][INFO][dist.py:818] - Using device='xpu' with backend='DDP' + 'ccl' for distributed training.\n    [2025-01-03 10:16:17.796091][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 0/47]\n    [2025-01-03 10:16:17.795365][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 2/47]\n    [2025-01-03 10:16:17.795355][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 3/47]\n    [2025-01-03 10:16:17.795344][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 5/47]\n    [2025-01-03 10:16:17.795374][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 1/47]\n    [2025-01-03 10:16:17.795396][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 4/47]\n    [2025-01-03 10:16:17.795352][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 6/47]\n    [2025-01-03 10:16:17.795331][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 7/47]\n    [2025-01-03 10:16:17.795358][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 8/47]\n    [2025-01-03 10:16:17.795331][INFO][dist.py:859] - ['x4509c0s2b0n0'][ 9/47]\n    [2025-01-03 10:16:17.795348][INFO][dist.py:859] - ['x4509c0s2b0n0'][10/47]\n    [2025-01-03 10:16:17.795332][INFO][dist.py:859] - ['x4509c0s2b0n0'][11/47]\n    [2025-01-03 10:16:17.795337][INFO][dist.py:859] - ['x4206c0s0b0n0'][20/47]\n    [2025-01-03 10:16:17.795322][INFO][dist.py:859] - ['x4206c0s0b0n0'][23/47]\n    [2025-01-03 10:16:17.795427][INFO][dist.py:859] - ['x4404c5s2b0n0'][36/47]\n    [2025-01-03 10:16:17.795368][INFO][dist.py:859] - ['x4206c0s0b0n0'][12/47]\n    [2025-01-03 10:16:17.795400][INFO][dist.py:859] - ['x4206c2s1b0n0'][24/47]\n    [2025-01-03 10:16:17.795402][INFO][dist.py:859] - ['x4404c5s2b0n0'][37/47]\n    [2025-01-03 10:16:17.795348][INFO][dist.py:859] - ['x4206c0s0b0n0'][13/47]\n    [2025-01-03 10:16:17.795375][INFO][dist.py:859] - ['x4206c2s1b0n0'][25/47]\n    [2025-01-03 10:16:17.795426][INFO][dist.py:859] - ['x4404c5s2b0n0'][38/47]\n    [2025-01-03 10:16:17.795356][INFO][dist.py:859] - ['x4206c0s0b0n0'][14/47]\n    [2025-01-03 10:16:17.795387][INFO][dist.py:859] - ['x4206c2s1b0n0'][26/47]\n    [2025-01-03 10:16:17.795350][INFO][dist.py:859] - ['x4206c0s0b0n0'][15/47]\n    [2025-01-03 10:16:17.795376][INFO][dist.py:859] - ['x4206c2s1b0n0'][27/47]\n    [2025-01-03 10:16:17.795391][INFO][dist.py:859] - ['x4206c0s0b0n0'][16/47]\n    [2025-01-03 10:16:17.795377][INFO][dist.py:859] - ['x4206c2s1b0n0'][28/47]\n    [2025-01-03 10:16:17.795341][INFO][dist.py:859] - ['x4206c0s0b0n0'][17/47]\n    [2025-01-03 10:16:17.795368][INFO][dist.py:859] - ['x4206c2s1b0n0'][29/47]\n    [2025-01-03 10:16:17.795429][INFO][dist.py:859] - ['x4206c0s0b0n0'][18/47]\n    [2025-01-03 10:16:17.795391][INFO][dist.py:859] - ['x4206c2s1b0n0'][30/47]\n    [2025-01-03 10:16:17.795406][INFO][dist.py:859] - ['x4404c5s2b0n0'][39/47]\n    [2025-01-03 10:16:17.795323][INFO][dist.py:859] - ['x4206c0s0b0n0'][19/47]\n    [2025-01-03 10:16:17.795362][INFO][dist.py:859] - ['x4206c2s1b0n0'][31/47]\n    [2025-01-03 10:16:17.795325][INFO][dist.py:859] - ['x4206c0s0b0n0'][21/47]\n    [2025-01-03 10:16:17.795381][INFO][dist.py:859] - ['x4206c2s1b0n0'][32/47]\n    [2025-01-03 10:16:17.795392][INFO][dist.py:859] - ['x4404c5s2b0n0'][43/47]\n    [2025-01-03 10:16:17.795338][INFO][dist.py:859] - ['x4206c0s0b0n0'][22/47]\n    [2025-01-03 10:16:17.795357][INFO][dist.py:859] - ['x4206c2s1b0n0'][33/47]\n    [2025-01-03 10:16:17.795403][INFO][dist.py:859] - ['x4404c5s2b0n0'][44/47]\n    [2025-01-03 10:16:17.795374][INFO][dist.py:859] - ['x4206c2s1b0n0'][34/47]\n    [2025-01-03 10:16:17.795396][INFO][dist.py:859] - ['x4404c5s2b0n0'][45/47]\n    [2025-01-03 10:16:17.795359][INFO][dist.py:859] - ['x4206c2s1b0n0'][35/47]\n    [2025-01-03 10:16:17.795402][INFO][dist.py:859] - ['x4404c5s2b0n0'][46/47]\n    [2025-01-03 10:16:17.795386][INFO][dist.py:859] - ['x4404c5s2b0n0'][47/47]\n    [2025-01-03 10:16:17.795442][INFO][dist.py:859] - ['x4404c5s2b0n0'][40/47]\n    [2025-01-03 10:16:17.795406][INFO][dist.py:859] - ['x4404c5s2b0n0'][41/47]\n    [2025-01-03 10:16:17.795542][INFO][dist.py:859] - ['x4404c5s2b0n0'][42/47]\n    [2025-01-03 10:16:18.127243][INFO][fsdp.py:168] -\n    =================================================================\n    Layer (type:depth-idx)                   Param #\n    =================================================================\n    Net                                      --\n    ├─Conv2d: 1-1                            320\n    ├─Conv2d: 1-2                            18,496\n    ├─Dropout: 1-3                           --\n    ├─Dropout: 1-4                           --\n    ├─Linear: 1-5                            1,179,776\n    ├─Linear: 1-6                            1,290\n    =================================================================\n    Total params: 1,199,882\n    Trainable params: 1,199,882\n    Non-trainable params: 0\n    =================================================================\n    [2025-01-03 10:16:18.158700][INFO][fsdp.py:187] - model=FullyShardedDataParallel(\n      (_fsdp_wrapped_module): Net(\n        (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))\n        (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))\n        (dropout1): Dropout(p=0.25, inplace=False)\n        (dropout2): Dropout(p=0.5, inplace=False)\n        (fc1): Linear(in_features=9216, out_features=128, bias=True)\n        (fc2): Linear(in_features=128, out_features=10, bias=True)\n      )\n    )\n    [2025-01-03 10:16:31.371926][INFO][fsdp.py:228] - epoch=1 dt=12.708439 train_loss=0.907240 test_loss=0.251744 test_acc=92.723282\n    [2025-01-03 10:16:32.432899][INFO][fsdp.py:228] - epoch=2 dt=0.552378 train_loss=0.295683 test_loss=0.146544 test_acc=95.474480\n    [2025-01-03 10:16:33.919076][INFO][fsdp.py:228] - epoch=3 dt=1.008603 train_loss=0.196732 test_loss=0.108416 test_acc=96.700554\n    [2025-01-03 10:16:35.221881][INFO][fsdp.py:228] - epoch=4 dt=0.497740 train_loss=0.162368 test_loss=0.088102 test_acc=97.288673\n    [2025-01-03 10:16:37.285276][INFO][fsdp.py:228] - epoch=5 dt=1.478215 train_loss=0.140347 test_loss=0.078072 test_acc=97.647530\n    [2025-01-03 10:16:38.332123][INFO][fsdp.py:228] - epoch=6 dt=0.535241 train_loss=0.129397 test_loss=0.072069 test_acc=97.836922\n    [2025-01-03 10:16:39.187470][INFO][fsdp.py:228] - epoch=7 dt=0.456952 train_loss=0.122222 test_loss=0.069216 test_acc=97.876793\n    [2025-01-03 10:16:40.155302][INFO][fsdp.py:228] - epoch=8 dt=0.508528 train_loss=0.113515 test_loss=0.066888 test_acc=97.956535\n    [2025-01-03 10:16:41.000475][INFO][fsdp.py:228] - epoch=9 dt=0.504310 train_loss=0.112428 test_loss=0.065184 test_acc=98.006378\n    [2025-01-03 10:16:41.867020][INFO][fsdp.py:228] - epoch=10 dt=0.513946 train_loss=0.110064 test_loss=0.064282 test_acc=98.026314\n    [2025-01-03 10:16:41.868205][INFO][fsdp.py:234] - 11 epochs took 23.7s\n    [2025-01-03 10:16:42.037421][INFO][history.py:723] - Saving epoch plot to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/plots/mplot\n    [2025-01-03 10:16:42.270616][INFO][history.py:723] - Saving dt plot to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/plots/mplot\n    [2025-01-03 10:16:42.489108][INFO][history.py:723] - Saving train_loss plot to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/plots/mplot\n    [2025-01-03 10:16:42.705195][INFO][history.py:723] - Saving test_loss plot to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/plots/mplot\n    [2025-01-03 10:16:42.916880][INFO][history.py:723] - Saving test_acc plot to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/plots/mplot\n    [2025-01-03 10:16:43.114870][INFO][history.py:603] - Saving tplots to /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/plots/tplot\n                               epoch [2025-01-03-101643]\n        ┌─────────────────────────────────────────────────────────────────────┐\n    10.0┤                                                                   ▄▞│\n        │                                                               ▄▄▀▀  │\n        │                                                           ▗▄▀▀      │\n     8.5┤                                                       ▗▄▞▀▘         │\n        │                                                   ▗▄▞▀▘             │\n        │                                               ▗▄▞▀▘                 │\n     7.0┤                                            ▄▞▀▘                     │\n        │                                        ▄▄▀▀                         │\n     5.5┤                                   ▗▄▄▀▀                             │\n        │                              ▗▄▄▀▀▘                                 │\n        │                           ▄▄▀▘                                      │\n     4.0┤                       ▄▄▀▀                                          │\n        │                   ▄▄▀▀                                              │\n        │               ▄▄▀▀                                                  │\n     2.5┤           ▗▄▞▀                                                      │\n        │       ▗▄▞▀▘                                                         │\n        │    ▄▄▀▘                                                             │\n     1.0┤▄▄▀▀                                                                 │\n        └┬───────┬──────┬───────┬──────┬───────┬──────┬───────┬──────┬────────┘\n         1       2      3       4      5       6      7       8      9\n    epoch\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/plots/tplot/epoch.txt\n                                dt [2025-01-03-101643]\n        ┌─────────────────────────────────────────────────────────────────────┐\n    12.7┤▌                                                                    │\n        │▚                                                                    │\n        │▝▖                                                                   │\n    10.7┤ ▚                                                                   │\n        │ ▐                                                                   │\n        │  ▌                                                                  │\n     8.6┤  ▐                                                                  │\n        │   ▌                                                                 │\n     6.6┤   ▚                                                                 │\n        │   ▝▖                                                                │\n        │    ▚                                                                │\n     4.5┤    ▐                                                                │\n        │     ▌                                                               │\n        │     ▐                                                               │\n     2.5┤      ▌                                                              │\n        │      ▚                                                              │\n        │      ▝▖       ▖            ▄▄▞▄▄▖                                   │\n     0.5┤       ▚▄▄▄▞▀▀▀▝▀▀▀▚▄▄▄▄▄▞▀▀     ▝▀▀▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄│\n        └┬───────┬──────┬───────┬──────┬───────┬──────┬───────┬──────┬────────┘\n         1       2      3       4      5       6      7       8      9\n    dt\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/plots/tplot/dt.txt\n                            train_loss [2025-01-03-101643]\n        ┌─────────────────────────────────────────────────────────────────────┐\n    0.91┤▌                                                                    │\n        │▐                                                                    │\n        │ ▌                                                                   │\n    0.77┤ ▐                                                                   │\n        │  ▚                                                                  │\n        │  ▝▖                                                                 │\n    0.64┤   ▚                                                                 │\n        │   ▝▖                                                                │\n    0.51┤    ▚                                                                │\n        │     ▌                                                               │\n        │     ▐                                                               │\n    0.38┤      ▌                                                              │\n        │      ▐                                                              │\n        │       ▚                                                             │\n    0.24┤        ▀▚▄▖                                                         │\n        │           ▝▀▚▄▖                                                     │\n        │               ▝▀▀▀▚▄▄▄▖                                             │\n    0.11┤                       ▝▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄│\n        └┬───────┬──────┬───────┬──────┬───────┬──────┬───────┬──────┬────────┘\n         1       2      3       4      5       6      7       8      9\n    train_loss\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/plots/tplot/train_loss.txt\n                              test_loss [2025-01-03-101643]\n         ┌────────────────────────────────────────────────────────────────────┐\n    0.252┤▌                                                                   │\n         │▝▖                                                                  │\n         │ ▐                                                                  │\n    0.221┤  ▚                                                                 │\n         │   ▌                                                                │\n         │   ▝▖                                                               │\n    0.189┤    ▐                                                               │\n         │     ▚                                                              │\n    0.158┤      ▌                                                             │\n         │      ▝▖                                                            │\n         │       ▝▄                                                           │\n    0.127┤         ▀▄                                                         │\n         │           ▀▄                                                       │\n         │             ▀▄▖                                                    │\n    0.096┤               ▝▀▄▄                                                 │\n         │                   ▀▀▄▄                                             │\n         │                       ▀▀▀▀▀▀▀▚▄▄▄                                  │\n    0.064┤                                  ▀▀▀▀▀▀▀▀▀▀▀▚▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄│\n         └┬──────┬───────┬──────┬───────┬──────┬───────┬──────┬───────┬───────┘\n          1      2       3      4       5      6       7      8       9\n    test_loss\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/plots/tplot/test_loss.txt\n                              test_acc [2025-01-03-101643]\n         ┌────────────────────────────────────────────────────────────────────┐\n    98.03┤                                     ▗▄▄▄▄▄▄▄▄▄▄▄▄▄▄▞▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀│\n         │                           ▗▄▄▀▀▀▀▀▀▀▘                              │\n         │                      ▗▄▄▀▀▘                                        │\n    97.14┤                   ▄▄▀▘                                             │\n         │               ▄▄▀▀                                                 │\n         │             ▗▞                                                     │\n    96.26┤           ▗▞▘                                                      │\n         │         ▗▞▘                                                        │\n    95.37┤       ▗▞▘                                                          │\n         │      ▗▘                                                            │\n         │     ▗▘                                                             │\n    94.49┤     ▌                                                              │\n         │    ▞                                                               │\n         │   ▞                                                                │\n    93.61┤  ▐                                                                 │\n         │ ▗▘                                                                 │\n         │▗▘                                                                  │\n    92.72┤▌                                                                   │\n         └┬──────┬───────┬──────┬───────┬──────┬───────┬──────┬───────┬───────┘\n          1      2       3      4       5      6       7      8       9\n    test_acc\n    text saved in /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/plots/tplot/test_acc.txt\n    [2025-01-03 10:16:43.287065][INFO][utils.py:132] - Saving dataset to: /lus/flare/projects/Aurora_deployment/foremans/projects/saforem2/mmm/outputs/mmm-fsdp/train_dataset.h5\n    [2025-01-03 10:16:43.333022][INFO][fsdp.py","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsaforem2%2Fmmm","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fsaforem2%2Fmmm","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsaforem2%2Fmmm/lists"}