Add files using upload-large-folder tool

100000/bf16_zero_pp_rank_12_mp_rank_00_optim_states.pt

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

100000/bf16_zero_pp_rank_13_mp_rank_00_optim_states.pt

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

100000/bf16_zero_pp_rank_14_mp_rank_00_optim_states.pt

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

100000/bf16_zero_pp_rank_16_mp_rank_00_optim_states.pt

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

100000/bf16_zero_pp_rank_18_mp_rank_00_optim_states.pt

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

100000/bf16_zero_pp_rank_24_mp_rank_00_optim_states.pt

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

100000/bf16_zero_pp_rank_28_mp_rank_00_optim_states.pt

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

100000/bf16_zero_pp_rank_30_mp_rank_00_optim_states.pt

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

100000/bf16_zero_pp_rank_31_mp_rank_00_optim_states.pt

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

100000/bf16_zero_pp_rank_4_mp_rank_00_optim_states.pt

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

100000/bf16_zero_pp_rank_6_mp_rank_00_optim_states.pt

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

100000/bf16_zero_pp_rank_7_mp_rank_00_optim_states.pt

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

100000/ema.pt.zst

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

100000/mp_rank_00_model_states.pt

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

100000/rng-0.ckpt

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

100000/rng-1.ckpt

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

100000/rng-10.ckpt

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

100000/rng-11.ckpt

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

100000/rng-12.ckpt

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

100000/rng-13.ckpt

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

100000/rng-14.ckpt

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

100000/rng-15.ckpt

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

100000/rng-16.ckpt

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

100000/rng-17.ckpt

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

100000/rng-18.ckpt

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

100000/rng-19.ckpt

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

100000/rng-2.ckpt

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

100000/rng-20.ckpt

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

100000/rng-21.ckpt

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

100000/rng-22.ckpt

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

100000/rng-23.ckpt

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

100000/rng-24.ckpt

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

100000/rng-25.ckpt

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

100000/rng-26.ckpt

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

100000/rng-27.ckpt

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

100000/rng-28.ckpt

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

100000/rng-29.ckpt

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

100000/rng-3.ckpt

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

100000/rng-30.ckpt

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

100000/rng-31.ckpt

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

100000/rng-4.ckpt

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

100000/rng-5.ckpt

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

100000/rng-6.ckpt

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

100000/rng-7.ckpt

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

100000/rng-8.ckpt

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

100000/rng-9.ckpt

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

config.json

@@ -0,0 +1,107 @@
+{
+  "attention": "self",
+  "base_config": {
+    "_name_or_path": "google/gemma-2b",
+    "add_cross_attention": false,
+    "architectures": [
+      "GemmaForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 2,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 1,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "head_dim": 128,
+    "hidden_act": "gelu",
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 2048,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 8192,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 8192,
+    "min_length": 0,
+    "model_type": "gemma",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 32,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 18,
+    "num_key_value_heads": 8,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 0,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 10000.0,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 256000
+  },
+  "dit_hidden_size": 2048,
+  "dit_num_hidden_layers": 30,
+  "in_channels": 16,
+  "initial_layers": 0,
+  "model_type": "DiT",
+  "out_channels": 16,
+  "patch_size": 2,
+  "pos_embed": "ape",
+  "pos_embed_max_size": 64,
+  "qk_norm": true,
+  "repa_enable": true,
+  "repa_enc_depth": 8,
+  "repa_projector_dim": 2048,
+  "repa_z_dim": 768,
+  "sample_size": 32,
+  "sandwich_norm": false,
+  "shared_attention_layers": "all",
+  "text_hidden_size": 2048,
+  "text_hidden_states_index": -1,
+  "text_modulation_embeds_dim": null,
+  "timestep_conditioning": "adaln-zero",
+  "transformers_version": "4.43.3"
+}

flops.txt

@@ -0,0 +1,934 @@
+
+-------------------------- DeepSpeed Flops Profiler --------------------------
+Profile Summary at step 2:
+Notations:
+data parallel size (dp_size), model parallel size(mp_size),
+number of parameters (params), number of multiply-accumulate operations(MACs),
+number of floating-point operations (flops), floating-point operations per second (FLOPS),
+fwd latency (forward propagation latency), bwd latency (backward propagation latency),
+step (weights update latency), iter latency (sum of fwd, bwd and step latency)
+
+world size:                                                             32      
+data parallel size:                                                     32      
+model parallel size:                                                    1       
+batch size per GPU:                                                     16      
+params per GPU:                                                         3.05 B  
+params of model = params per GPU * mp_size:                             3.05 B  
+fwd MACs per GPU:                                                       9.86 TMACs
+fwd flops per GPU:                                                      19.73 T 
+fwd flops of model = fwd flops per GPU * mp_size:                       19.73 T 
+fwd latency:                                                            158.91 ms
+fwd FLOPS per GPU = fwd flops per GPU / fwd latency:                    124.15 TFLOPS
+bwd latency:                                                            559.86 ms
+bwd FLOPS per GPU = 2 * fwd flops per GPU / bwd latency:                70.48 TFLOPS
+fwd+bwd FLOPS per GPU = 3 * fwd flops per GPU / (fwd+bwd latency):      82.34 TFLOPS
+step latency:                                                           167.69 ms
+iter latency:                                                           886.46 ms
+FLOPS per GPU = 3 * fwd flops per GPU / iter latency:                   66.77 TFLOPS
+samples/second:                                                         577.58  
+
+----------------------------- Aggregated Profile per GPU -----------------------------
+Top 1 modules in terms of params, MACs or fwd latency at different model depths:
+depth 0:
+    params      - {'DiT': '3.05 B'}
+    MACs        - {'DiT': '9.86 TMACs'}
+    fwd latency - {'DiT': '158.73 ms'}
+depth 1:
+    params      - {'ModuleList': '3.02 B'}
+    MACs        - {'ModuleList': '9.8 TMACs'}
+    fwd latency - {'ModuleList': '154.18 ms'}
+depth 2:
+    params      - {'DiTLayer': '3.02 B'}
+    MACs        - {'DiTLayer': '9.8 TMACs'}
+    fwd latency - {'DiTLayer': '154.18 ms'}
+depth 3:
+    params      - {'GemmaMLP': '1.51 B'}
+    MACs        - {'GemmaMLP': '6.18 TMACs'}
+    fwd latency - {'DiTSelfAttention': '81.24 ms'}
+
+------------------------------ Detailed Profile per GPU ------------------------------
+Each module profile is listed after its name in the following order: 
+params, percentage of total params, MACs, percentage of total MACs, fwd latency, percentage of total fwd latency, fwd FLOPS
+
+Note: 1. A module can have torch.nn.module or torch.nn.functional to compute logits (e.g. CrossEntropyLoss). They are not counted as submodules, thus not to be printed out. However they make up the difference between a parent's MACs (or latency) and the sum of its submodules'.
+2. Number of floating-point operations is a theoretical estimation, thus FLOPS computed using that could be larger than the maximum system throughput.
+3. The fwd latency listed in the top module's profile is directly captured at the module forward function in PyTorch, thus it's less than the fwd latency shown above which is captured in DeepSpeed.
+
+DiT(
+  3.05 B = 100% Params, 9.86 TMACs = 100% MACs, 158.73 ms = 100% latency, 124.29 TFLOPS
+  (layers): ModuleList(
+    (0): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.24 ms = 3.3% latency, 124.79 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 628.95 us = 0.4% latency, 1.28 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 36.24 us = 0.02% latency, 904.2 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 184.54 us = 0.12% latency, 4.36 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 237.46 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.76 ms = 1.74% latency, 87.27 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.77 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 238.42 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 276.8 us = 0.17% latency, 248.26 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 160.93 us = 0.1% latency, 106.75 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 151.87 us = 0.1% latency, 113.12 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 137.09 us = 0.09% latency, 125.32 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 136.85 us = 0.09% latency, 125.54 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 251.05 us = 0.16% latency, 273.72 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.18 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.23 ms = 0.77% latency, 336.48 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 343.8 us = 0.22% latency, 399.76 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 322.1 us = 0.2% latency, 426.69 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 301.36 us = 0.19% latency, 456.06 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 86.55 us = 0.05% latency, 387.71 GFLOPS)
+      )
+    )
+    (1): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.12 ms = 3.23% latency, 127.56 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 585.56 us = 0.37% latency, 1.38 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 33.38 us = 0.02% latency, 981.71 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 159.98 us = 0.1% latency, 5.03 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.56 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.71 ms = 1.71% latency, 88.76 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.77 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 248.19 us = 0.16% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 248.67 us = 0.16% latency, 276.35 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 156.64 us = 0.1% latency, 109.68 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 149.73 us = 0.09% latency, 114.74 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 139.47 us = 0.09% latency, 123.18 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 136.61 us = 0.09% latency, 125.75 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 242.95 us = 0.15% latency, 282.86 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 237.94 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.21 ms = 0.76% latency, 341.33 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 334.02 us = 0.21% latency, 411.46 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 319.96 us = 0.2% latency, 429.55 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 301.36 us = 0.19% latency, 456.06 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 83.68 us = 0.05% latency, 400.96 GFLOPS)
+      )
+    )
+    (2): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.12 ms = 3.23% latency, 127.59 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 585.08 us = 0.37% latency, 1.38 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 34.33 us = 0.02% latency, 954.44 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 160.69 us = 0.1% latency, 5.01 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.66 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.71 ms = 1.71% latency, 88.83 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.06 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 235.08 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 254.15 us = 0.16% latency, 270.38 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 154.26 us = 0.1% latency, 111.37 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 159.03 us = 0.1% latency, 108.03 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 139.47 us = 0.09% latency, 123.18 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 136.38 us = 0.09% latency, 125.97 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 238.9 us = 0.15% latency, 287.66 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 236.03 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.21 ms = 0.76% latency, 340.46 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 333.55 us = 0.21% latency, 412.05 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 322.82 us = 0.2% latency, 425.75 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 299.93 us = 0.19% latency, 458.24 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 85.59 us = 0.05% latency, 392.03 GFLOPS)
+      )
+    )
+    (3): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.11 ms = 3.22% latency, 127.97 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 602.25 us = 0.38% latency, 1.34 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 33.14 us = 0.02% latency, 988.77 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 164.27 us = 0.1% latency, 4.9 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 241.04 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.68 ms = 1.69% latency, 89.59 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.29 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 234.84 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 247.48 us = 0.16% latency, 277.68 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 154.26 us = 0.1% latency, 111.37 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 151.16 us = 0.1% latency, 113.66 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 139.24 us = 0.09% latency, 123.39 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 136.85 us = 0.09% latency, 125.54 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 233.41 us = 0.15% latency, 294.41 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.66 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.2 ms = 0.76% latency, 342.82 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 331.88 us = 0.21% latency, 414.12 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 317.81 us = 0.2% latency, 432.45 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 300.41 us = 0.19% latency, 457.51 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 82.49 us = 0.05% latency, 406.76 GFLOPS)
+      )
+    )
+    (4): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.09 ms = 3.21% latency, 128.43 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 581.98 us = 0.37% latency, 1.38 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 33.14 us = 0.02% latency, 988.77 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 159.98 us = 0.1% latency, 5.03 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 241.04 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.69 ms = 1.69% latency, 89.53 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.53 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 234.13 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 247.24 us = 0.16% latency, 277.95 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 157.36 us = 0.1% latency, 109.18 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 151.63 us = 0.1% latency, 113.3 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 139.47 us = 0.09% latency, 123.18 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 137.57 us = 0.09% latency, 124.88 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 237.46 us = 0.15% latency, 289.39 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.61 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.2 ms = 0.76% latency, 343.16 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 329.73 us = 0.21% latency, 416.82 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 317.81 us = 0.2% latency, 432.45 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 302.55 us = 0.19% latency, 454.26 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 82.25 us = 0.05% latency, 407.93 GFLOPS)
+      )
+    )
+    (5): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.09 ms = 3.21% latency, 128.42 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 581.5 us = 0.37% latency, 1.38 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 34.33 us = 0.02% latency, 954.44 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 158.55 us = 0.1% latency, 5.08 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.56 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.69 ms = 1.69% latency, 89.58 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 461.58 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 234.84 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 242.23 us = 0.15% latency, 283.69 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 156.4 us = 0.1% latency, 109.84 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 151.16 us = 0.1% latency, 113.66 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 141.62 us = 0.09% latency, 121.31 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 136.14 us = 0.09% latency, 126.2 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 239.85 us = 0.15% latency, 286.51 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.37 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.2 ms = 0.76% latency, 342.68 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 335.93 us = 0.21% latency, 409.13 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 320.2 us = 0.2% latency, 429.23 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 297.78 us = 0.19% latency, 461.54 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 81.06 us = 0.05% latency, 413.93 GFLOPS)
+      )
+    )
+    (6): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.16 ms = 3.25% latency, 126.68 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 582.22 us = 0.37% latency, 1.38 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 33.38 us = 0.02% latency, 981.71 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 159.03 us = 0.1% latency, 5.06 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.09 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.72 ms = 1.71% latency, 88.51 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.53 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 235.56 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 249.86 us = 0.16% latency, 275.03 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 159.98 us = 0.1% latency, 107.39 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 152.11 us = 0.1% latency, 112.94 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 143.05 us = 0.09% latency, 120.1 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 139.24 us = 0.09% latency, 123.39 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 236.99 us = 0.15% latency, 289.97 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.66 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.23 ms = 0.78% latency, 334.53 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 336.65 us = 0.21% latency, 408.26 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 330.69 us = 0.21% latency, 415.62 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 304.7 us = 0.19% latency, 451.06 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 85.59 us = 0.05% latency, 392.03 GFLOPS)
+      )
+    )
+    (7): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.19 ms = 3.27% latency, 126.06 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 600.58 us = 0.38% latency, 1.34 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 34.81 us = 0.02% latency, 941.36 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 169.99 us = 0.11% latency, 4.74 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 241.04 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.72 ms = 1.71% latency, 88.39 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 463.01 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 235.08 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 257.25 us = 0.16% latency, 267.13 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 156.16 us = 0.1% latency, 110.01 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 150.68 us = 0.09% latency, 114.02 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 139.95 us = 0.09% latency, 122.76 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 140.91 us = 0.09% latency, 121.92 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 247.72 us = 0.16% latency, 277.41 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.61 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.23 ms = 0.78% latency, 334.72 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 340.7 us = 0.21% latency, 403.4 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 321.63 us = 0.2% latency, 427.32 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 311.14 us = 0.2% latency, 441.73 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 83.45 us = 0.05% latency, 402.11 GFLOPS)
+      )
+    )
+    (8): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.15 ms = 3.24% latency, 127.04 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 592.95 us = 0.37% latency, 1.36 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 34.09 us = 0.02% latency, 961.11 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 154.02 us = 0.1% latency, 5.23 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 241.04 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.72 ms = 1.71% latency, 88.56 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.53 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 235.08 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 249.62 us = 0.16% latency, 275.29 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 153.06 us = 0.1% latency, 112.24 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 149.73 us = 0.09% latency, 114.74 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 141.86 us = 0.09% latency, 121.11 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 148.53 us = 0.09% latency, 115.66 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 244.62 us = 0.15% latency, 280.93 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.18 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.22 ms = 0.77% latency, 339.19 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 334.98 us = 0.21% latency, 410.29 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 318.29 us = 0.2% latency, 431.81 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 306.13 us = 0.19% latency, 448.96 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 83.21 us = 0.05% latency, 403.26 GFLOPS)
+      )
+    )
+    (9): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.17 ms = 3.26% latency, 126.46 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 590.8 us = 0.37% latency, 1.36 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 33.86 us = 0.02% latency, 967.88 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 161.17 us = 0.1% latency, 5 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 241.76 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.72 ms = 1.71% latency, 88.44 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 463.01 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 235.08 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 255.11 us = 0.16% latency, 269.37 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 161.17 us = 0.1% latency, 106.59 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 148.53 us = 0.09% latency, 115.66 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 140.67 us = 0.09% latency, 122.13 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 137.57 us = 0.09% latency, 124.88 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 255.58 us = 0.16% latency, 268.87 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 236.75 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.23 ms = 0.78% latency, 334.14 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 337.84 us = 0.21% latency, 406.82 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 338.55 us = 0.21% latency, 405.96 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 300.17 us = 0.19% latency, 457.87 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 83.45 us = 0.05% latency, 402.11 GFLOPS)
+      )
+    )
+    (10): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.11 ms = 3.22% latency, 127.93 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 581.98 us = 0.37% latency, 1.38 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 33.86 us = 0.02% latency, 967.88 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 161.89 us = 0.1% latency, 4.97 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.66 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.69 ms = 1.7% latency, 89.31 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 461.58 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 237.94 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 244.62 us = 0.15% latency, 280.93 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 157.12 us = 0.1% latency, 109.34 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 150.68 us = 0.09% latency, 114.02 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 137.33 us = 0.09% latency, 125.1 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 135.9 us = 0.09% latency, 126.42 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 247.48 us = 0.16% latency, 277.68 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 237.94 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.21 ms = 0.76% latency, 340.39 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 336.17 us = 0.21% latency, 408.84 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 319.48 us = 0.2% latency, 430.19 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 301.12 us = 0.19% latency, 456.42 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 81.06 us = 0.05% latency, 413.93 GFLOPS)
+      )
+    )
+    (11): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.07 ms = 3.19% latency, 128.95 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 577.21 us = 0.36% latency, 1.4 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 32.66 us = 0.02% latency, 1 GFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 159.03 us = 0.1% latency, 5.06 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 237.94 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.67 ms = 1.68% latency, 90.06 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 461.1 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 234.6 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 244.14 us = 0.15% latency, 281.47 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 152.59 us = 0.1% latency, 112.59 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 149.97 us = 0.09% latency, 114.56 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 141.14 us = 0.09% latency, 121.72 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 134.71 us = 0.08% latency, 127.54 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 235.32 us = 0.15% latency, 292.03 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 236.99 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.2 ms = 0.76% latency, 342.48 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 334.98 us = 0.21% latency, 410.29 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 317.57 us = 0.2% latency, 432.78 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 300.17 us = 0.19% latency, 457.87 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 80.82 us = 0.05% latency, 415.15 GFLOPS)
+      )
+    )
+    (12): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.11 ms = 3.22% latency, 128.03 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 587.22 us = 0.37% latency, 1.37 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 34.33 us = 0.02% latency, 954.44 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 163.32 us = 0.1% latency, 4.93 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.42 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.7 ms = 1.7% latency, 89.2 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.29 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 233.41 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 245.33 us = 0.15% latency, 280.11 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 159.26 us = 0.1% latency, 107.87 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 151.16 us = 0.1% latency, 113.66 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 140.91 us = 0.09% latency, 121.92 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 136.38 us = 0.09% latency, 125.97 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 235.32 us = 0.15% latency, 292.03 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.09 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.2 ms = 0.76% latency, 343.16 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 332.12 us = 0.21% latency, 413.83 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 317.57 us = 0.2% latency, 432.78 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 300.41 us = 0.19% latency, 457.51 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 81.54 us = 0.05% latency, 411.51 GFLOPS)
+      )
+    )
+    (13): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.11 ms = 3.22% latency, 127.94 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 582.7 us = 0.37% latency, 1.38 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 32.19 us = 0.02% latency, 1.02 GFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 158.79 us = 0.1% latency, 5.07 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.56 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.68 ms = 1.69% latency, 89.82 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.53 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 232.93 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 248.43 us = 0.16% latency, 276.61 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 154.73 us = 0.1% latency, 111.03 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 150.68 us = 0.09% latency, 114.02 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 138.04 us = 0.09% latency, 124.45 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 134.94 us = 0.09% latency, 127.31 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 234.84 us = 0.15% latency, 292.62 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 236.99 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.23 ms = 0.78% latency, 334.98 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 349.76 us = 0.22% latency, 392.95 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 322.82 us = 0.2% latency, 425.75 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 302.79 us = 0.19% latency, 453.91 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 82.73 us = 0.05% latency, 405.58 GFLOPS)
+      )
+    )
+    (14): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.11 ms = 3.22% latency, 127.87 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 584.13 us = 0.37% latency, 1.38 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 33.38 us = 0.02% latency, 981.71 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 160.69 us = 0.1% latency, 5.01 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.09 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.69 ms = 1.7% latency, 89.25 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 461.34 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 244.14 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 244.62 us = 0.15% latency, 280.93 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 154.97 us = 0.1% latency, 110.86 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 148.53 us = 0.09% latency, 115.66 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 138.04 us = 0.09% latency, 124.45 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 135.9 us = 0.09% latency, 126.42 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 241.52 us = 0.15% latency, 284.53 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.33 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.21 ms = 0.76% latency, 340.99 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 337.12 us = 0.21% latency, 407.68 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 319.72 us = 0.2% latency, 429.87 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 299.45 us = 0.19% latency, 458.97 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 81.3 us = 0.05% latency, 412.72 GFLOPS)
+      )
+    )
+    (15): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.11 ms = 3.22% latency, 127.93 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 578.88 us = 0.36% latency, 1.39 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 32.66 us = 0.02% latency, 1 GFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 158.55 us = 0.1% latency, 5.08 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.33 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.71 ms = 1.7% latency, 88.87 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 463.25 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 231.74 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 246.29 us = 0.16% latency, 279.02 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 155.45 us = 0.1% latency, 110.52 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 150.68 us = 0.09% latency, 114.02 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 141.86 us = 0.09% latency, 121.11 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 150.44 us = 0.09% latency, 114.2 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 239.37 us = 0.15% latency, 287.08 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 237.46 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.21 ms = 0.76% latency, 342.07 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 335.22 us = 0.21% latency, 410 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 319 us = 0.2% latency, 430.84 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 298.74 us = 0.19% latency, 460.06 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 81.3 us = 0.05% latency, 412.72 GFLOPS)
+      )
+    )
+    (16): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.1 ms = 3.22% latency, 128.09 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 578.64 us = 0.36% latency, 1.39 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 33.14 us = 0.02% latency, 988.77 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 159.26 us = 0.1% latency, 5.06 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.9 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.69 ms = 1.69% latency, 89.54 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.06 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 231.27 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 245.33 us = 0.15% latency, 280.11 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 153.54 us = 0.1% latency, 111.89 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 147.58 us = 0.09% latency, 116.41 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 140.19 us = 0.09% latency, 122.55 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 136.85 us = 0.09% latency, 125.54 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 244.14 us = 0.15% latency, 281.47 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 237.46 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.22 ms = 0.77% latency, 339.12 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 338.79 us = 0.21% latency, 405.67 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 319.96 us = 0.2% latency, 429.55 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 301.6 us = 0.19% latency, 455.7 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 81.54 us = 0.05% latency, 411.51 GFLOPS)
+      )
+    )
+    (17): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.1 ms = 3.21% latency, 128.11 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 573.16 us = 0.36% latency, 1.41 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 32.9 us = 0.02% latency, 995.93 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 156.88 us = 0.1% latency, 5.13 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 236.75 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.69 ms = 1.69% latency, 89.43 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 461.82 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 234.13 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 247.72 us = 0.16% latency, 277.41 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 157.59 us = 0.1% latency, 109.01 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 148.3 us = 0.09% latency, 115.85 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 139.24 us = 0.09% latency, 123.39 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 135.66 us = 0.09% latency, 126.64 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 242.47 us = 0.15% latency, 283.41 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.13 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.22 ms = 0.77% latency, 339.06 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 339.03 us = 0.21% latency, 405.39 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 320.67 us = 0.2% latency, 428.6 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 301.6 us = 0.19% latency, 455.7 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 82.02 us = 0.05% latency, 409.12 GFLOPS)
+      )
+    )
+    (18): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.19 ms = 3.27% latency, 125.93 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 584.6 us = 0.37% latency, 1.38 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 33.14 us = 0.02% latency, 988.77 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 160.22 us = 0.1% latency, 5.03 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 241.76 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.74 ms = 1.72% latency, 87.86 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 469.68 us = 0.3% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 239.85 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 245.33 us = 0.15% latency, 280.11 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 158.79 us = 0.1% latency, 108.19 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 153.78 us = 0.1% latency, 111.72 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 143.77 us = 0.09% latency, 119.5 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 140.91 us = 0.09% latency, 121.92 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 249.86 us = 0.16% latency, 275.03 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.13 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.24 ms = 0.78% latency, 332.22 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 349.28 us = 0.22% latency, 393.49 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 328.06 us = 0.21% latency, 418.94 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 306.61 us = 0.19% latency, 448.26 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 83.92 us = 0.05% latency, 399.82 GFLOPS)
+      )
+    )
+    (19): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.26 ms = 3.32% latency, 124.18 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 604.15 us = 0.38% latency, 1.33 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 37.67 us = 0.02% latency, 869.87 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 167.61 us = 0.11% latency, 4.8 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.33 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.78 ms = 1.75% latency, 86.61 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 465.63 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 250.34 us = 0.16% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 255.11 us = 0.16% latency, 269.37 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 166.18 us = 0.1% latency, 103.38 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 156.88 us = 0.1% latency, 109.51 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 149.01 us = 0.09% latency, 115.29 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 141.38 us = 0.09% latency, 121.51 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 253.68 us = 0.16% latency, 270.89 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.37 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.25 ms = 0.79% latency, 330.25 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 349.76 us = 0.22% latency, 392.95 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 326.4 us = 0.21% latency, 421.08 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 311.61 us = 0.2% latency, 441.06 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 83.92 us = 0.05% latency, 399.82 GFLOPS)
+      )
+    )
+    (20): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.13 ms = 3.23% latency, 127.45 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 586.75 us = 0.37% latency, 1.37 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 34.57 us = 0.02% latency, 947.85 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 161.89 us = 0.1% latency, 4.97 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.13 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.69 ms = 1.69% latency, 89.53 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 463.25 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 232.93 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 243.43 us = 0.15% latency, 282.3 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 154.97 us = 0.1% latency, 110.86 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 147.1 us = 0.09% latency, 116.79 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 142.57 us = 0.09% latency, 120.5 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 135.9 us = 0.09% latency, 126.42 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 241.99 us = 0.15% latency, 283.97 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.33 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.23 ms = 0.77% latency, 335.37 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 342.37 us = 0.22% latency, 401.44 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 322.82 us = 0.2% latency, 425.75 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 307.56 us = 0.19% latency, 446.87 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 82.97 us = 0.05% latency, 404.42 GFLOPS)
+      )
+    )
+    (21): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.15 ms = 3.25% latency, 126.85 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 586.03 us = 0.37% latency, 1.37 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 34.09 us = 0.02% latency, 961.11 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 160.22 us = 0.1% latency, 5.03 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.33 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.73 ms = 1.72% latency, 88.03 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 463.72 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 235.08 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 256.06 us = 0.16% latency, 268.37 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 163.79 us = 0.1% latency, 104.89 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 151.87 us = 0.1% latency, 113.12 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 141.62 us = 0.09% latency, 121.31 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 136.61 us = 0.09% latency, 125.75 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 248.67 us = 0.16% latency, 276.35 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.9 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.21 ms = 0.76% latency, 340.59 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 337.6 us = 0.21% latency, 407.11 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 320.67 us = 0.2% latency, 428.6 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 300.41 us = 0.19% latency, 457.51 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 81.78 us = 0.05% latency, 410.31 GFLOPS)
+      )
+    )
+    (22): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.14 ms = 3.24% latency, 127.25 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 585.79 us = 0.37% latency, 1.37 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 34.09 us = 0.02% latency, 961.11 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 163.08 us = 0.1% latency, 4.94 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 241.52 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.71 ms = 1.71% latency, 88.84 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.77 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 239.37 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 258.21 us = 0.16% latency, 266.14 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 154.5 us = 0.1% latency, 111.2 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 149.01 us = 0.09% latency, 115.29 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 139.95 us = 0.09% latency, 122.76 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 137.33 us = 0.09% latency, 125.1 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 237.94 us = 0.15% latency, 288.81 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.37 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.22 ms = 0.77% latency, 337.86 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 340.22 us = 0.21% latency, 403.97 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 323.53 us = 0.2% latency, 424.81 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 301.12 us = 0.19% latency, 456.42 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 82.25 us = 0.05% latency, 407.93 GFLOPS)
+      )
+    )
+    (23): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.15 ms = 3.25% latency, 126.88 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 595.81 us = 0.38% latency, 1.35 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 34.81 us = 0.02% latency, 941.36 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 169.52 us = 0.11% latency, 4.75 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.61 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.71 ms = 1.71% latency, 88.66 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 466.35 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 238.9 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 249.15 us = 0.16% latency, 275.82 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 153.54 us = 0.1% latency, 111.89 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 150.2 us = 0.09% latency, 114.38 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 142.34 us = 0.09% latency, 120.7 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 137.09 us = 0.09% latency, 125.32 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 243.66 us = 0.15% latency, 282.03 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 236.27 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.22 ms = 0.77% latency, 337.34 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 342.61 us = 0.22% latency, 401.16 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 320.43 us = 0.2% latency, 428.91 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 303.27 us = 0.19% latency, 453.19 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 82.73 us = 0.05% latency, 405.58 GFLOPS)
+      )
+    )
+    (24): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.12 ms = 3.23% latency, 127.67 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 581.98 us = 0.37% latency, 1.38 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 35.29 us = 0.02% latency, 928.64 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 160.22 us = 0.1% latency, 5.03 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.66 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.69 ms = 1.7% latency, 89.29 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 461.34 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 235.08 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 246.76 us = 0.16% latency, 278.48 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 157.83 us = 0.1% latency, 108.85 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 149.01 us = 0.09% latency, 115.29 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 140.67 us = 0.09% latency, 122.13 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 136.85 us = 0.09% latency, 125.54 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 244.86 us = 0.15% latency, 280.65 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.13 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.22 ms = 0.77% latency, 337.6 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 340.22 us = 0.21% latency, 403.97 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 326.16 us = 0.21% latency, 421.39 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 303.51 us = 0.19% latency, 452.84 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 82.25 us = 0.05% latency, 407.93 GFLOPS)
+      )
+    )
+    (25): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.12 ms = 3.22% latency, 127.78 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 585.56 us = 0.37% latency, 1.38 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 33.86 us = 0.02% latency, 967.88 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 159.74 us = 0.1% latency, 5.04 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 241.04 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.69 ms = 1.7% latency, 89.35 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 462.77 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 235.08 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 246.29 us = 0.16% latency, 279.02 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 155.21 us = 0.1% latency, 110.69 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 151.87 us = 0.1% latency, 113.12 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 141.38 us = 0.09% latency, 121.51 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 137.57 us = 0.09% latency, 124.88 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 242.95 us = 0.15% latency, 282.86 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.13 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.21 ms = 0.76% latency, 339.86 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 337.12 us = 0.21% latency, 407.68 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 321.63 us = 0.2% latency, 427.32 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 300.88 us = 0.19% latency, 456.78 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 81.78 us = 0.05% latency, 410.31 GFLOPS)
+      )
+    )
+    (26): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.18 ms = 3.26% latency, 126.29 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 589.61 us = 0.37% latency, 1.37 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 33.62 us = 0.02% latency, 974.74 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 165.22 us = 0.1% latency, 4.87 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 239.13 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.71 ms = 1.71% latency, 88.78 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 461.82 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 239.61 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 250.1 us = 0.16% latency, 274.77 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 155.69 us = 0.1% latency, 110.35 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 149.73 us = 0.09% latency, 114.74 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 139.47 us = 0.09% latency, 123.18 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 135.9 us = 0.09% latency, 126.42 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 250.34 us = 0.16% latency, 274.51 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 240.09 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.24 ms = 0.78% latency, 332.92 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 346.66 us = 0.22% latency, 396.47 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 329.02 us = 0.21% latency, 417.73 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 307.32 us = 0.19% latency, 447.22 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 83.45 us = 0.05% latency, 402.11 GFLOPS)
+      )
+    )
+    (27): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.15 ms = 3.25% latency, 126.81 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 590.09 us = 0.37% latency, 1.36 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 34.57 us = 0.02% latency, 947.85 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 165.22 us = 0.1% latency, 4.87 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 237.7 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.71 ms = 1.71% latency, 88.7 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 463.25 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 234.13 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 252.25 us = 0.16% latency, 272.43 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 160.46 us = 0.1% latency, 107.07 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 151.4 us = 0.1% latency, 113.48 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 141.38 us = 0.09% latency, 121.51 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 139.24 us = 0.09% latency, 123.39 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 243.66 us = 0.15% latency, 282.03 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.66 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.23 ms = 0.77% latency, 336.03 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 341.65 us = 0.22% latency, 402.28 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 325.44 us = 0.21% latency, 422.32 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 304.7 us = 0.19% latency, 451.06 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 83.21 us = 0.05% latency, 403.26 GFLOPS)
+      )
+    )
+    (28): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.17 ms = 3.26% latency, 126.45 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 592.71 us = 0.37% latency, 1.36 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 35.29 us = 0.02% latency, 928.64 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 166.18 us = 0.1% latency, 4.85 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.66 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.73 ms = 1.72% latency, 88.07 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 464.92 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 245.33 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 252.96 us = 0.16% latency, 271.66 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 158.55 us = 0.1% latency, 108.36 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 157.12 us = 0.1% latency, 109.34 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 143.53 us = 0.09% latency, 119.7 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 137.57 us = 0.09% latency, 124.88 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 241.04 us = 0.15% latency, 285.09 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 237.46 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.22 ms = 0.77% latency, 337.8 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 342.37 us = 0.22% latency, 401.44 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 321.87 us = 0.2% latency, 427.01 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 302.55 us = 0.19% latency, 454.26 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 82.49 us = 0.05% latency, 406.76 GFLOPS)
+      )
+    )
+    (29): DiTLayer(
+      100.68 M = 3.3% Params, 326.82 GMACs = 3.31% MACs, 5.16 ms = 3.25% latency, 126.68 TFLOPS
+      (input_layernorm): AdaLayerNormZero(
+        25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 577.21 us = 0.36% latency, 1.4 TFLOPS
+        (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 35.05 us = 0.02% latency, 934.96 MFLOPS)
+        (linear): Linear(25.18 M = 0.83% Params, 402.65 MMACs = 0% MACs, 158.79 us = 0.1% latency, 5.07 TFLOPS, in_features=2048, out_features=12288, bias=True)
+        (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 236.75 us = 0.15% latency, 0 FLOPS)
+      )
+      (self_attn): DiTSelfAttention(
+        25.17 M = 0.83% Params, 120.26 GMACs = 1.22% MACs, 2.73 ms = 1.72% latency, 88 TFLOPS
+        (q_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 463.72 us = 0.29% latency, 0 FLOPS)
+        (k_norm): GemmaRMSNorm(128 = 0% Params, 0 MACs = 0% MACs, 237.94 us = 0.15% latency, 0 FLOPS)
+        (q_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 247.48 us = 0.16% latency, 277.68 TFLOPS, in_features=2048, out_features=4096, bias=False)
+        (k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 165.94 us = 0.1% latency, 103.53 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 155.21 us = 0.1% latency, 110.69 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_k_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 141.14 us = 0.09% latency, 121.72 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (text_v_proj): Linear(2.1 M = 0.07% Params, 8.59 GMACs = 0.09% MACs, 139.95 us = 0.09% latency, 122.76 TFLOPS, in_features=2048, out_features=1024, bias=False)
+        (o_proj): Linear(8.39 M = 0.28% Params, 34.36 GMACs = 0.35% MACs, 252.25 us = 0.16% latency, 272.43 TFLOPS, in_features=4096, out_features=2048, bias=False)
+      )
+      (post_attention_layernorm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 237.7 us = 0.15% latency, 0 FLOPS)
+      (mlp): GemmaMLP(
+        50.33 M = 1.65% Params, 206.16 GMACs = 2.09% MACs, 1.23 ms = 0.77% latency, 336.09 TFLOPS
+        (gate_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 341.42 us = 0.22% latency, 402.56 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (up_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 321.15 us = 0.2% latency, 427.96 TFLOPS, in_features=2048, out_features=8192, bias=False)
+        (down_proj): Linear(16.78 M = 0.55% Params, 68.72 GMACs = 0.7% MACs, 301.84 us = 0.19% latency, 455.34 TFLOPS, in_features=8192, out_features=2048, bias=False)
+        (act_fn): PytorchGELUTanh(0 = 0% Params, 0 MACs = 0% MACs, 82.49 us = 0.05% latency, 406.76 GFLOPS)
+      )
+    )
+  )
+  (patch_embed): PatchEmbed(
+    133.12 K = 0% Params, 536.87 MMACs = 0.01% MACs, 487.8 us = 0.31% latency, 2.22 TFLOPS
+    (proj): Conv2d(133.12 K = 0% Params, 536.87 MMACs = 0.01% MACs, 300.17 us = 0.19% latency, 3.61 TFLOPS, 16, 2048, kernel_size=(2, 2), stride=(2, 2))
+  )
+  (rotary_emb): GemmaRotaryEmbedding(0 = 0% Params, 0 MACs = 0% MACs, 0 s = 0% latency, 0 FLOPS)
+  (time_proj): Timesteps(0 = 0% Params, 0 MACs = 0% MACs, 248.43 us = 0.16% latency, 0 FLOPS)
+  (timestep_embedder): Sequential(
+    4.72 M = 0.15% Params, 75.5 MMACs = 0% MACs, 493.53 us = 0.31% latency, 306.02 GFLOPS
+    (0): Linear(526.34 K = 0.02% Params, 8.39 MMACs = 0% MACs, 227.93 us = 0.14% latency, 73.61 GFLOPS, in_features=256, out_features=2048, bias=True)
+    (1): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 50.78 us = 0.03% latency, 645.25 MFLOPS)
+    (2): Linear(4.2 M = 0.14% Params, 67.11 MMACs = 0% MACs, 144.72 us = 0.09% latency, 927.43 GFLOPS, in_features=2048, out_features=2048, bias=True)
+  )
+  (context_embedder): Sequential(
+    4.2 M = 0.14% Params, 17.18 GMACs = 0.17% MACs, 451.56 us = 0.28% latency, 76.09 TFLOPS
+    (0): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 171.18 us = 0.11% latency, 0 FLOPS)
+    (1): Linear(4.2 M = 0.14% Params, 17.18 GMACs = 0.17% MACs, 226.97 us = 0.14% latency, 151.38 TFLOPS, in_features=2048, out_features=2048, bias=True)
+  )
+  (norm_out): AdaLayerNormOut(
+    8.39 M = 0.28% Params, 134.22 MMACs = 0% MACs, 570.77 us = 0.36% latency, 470.36 GFLOPS
+    (silu): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 35.05 us = 0.02% latency, 934.96 MFLOPS)
+    (linear): Linear(8.39 M = 0.28% Params, 134.22 MMACs = 0% MACs, 148.77 us = 0.09% latency, 1.8 TFLOPS, in_features=2048, out_features=4096, bias=True)
+    (norm): GemmaRMSNorm(2.05 K = 0% Params, 0 MACs = 0% MACs, 238.66 us = 0.15% latency, 0 FLOPS)
+  )
+  (proj_out): Linear(131.14 K = 0% Params, 536.87 MMACs = 0.01% MACs, 173.09 us = 0.11% latency, 6.2 TFLOPS, in_features=2048, out_features=64, bias=True)
+  (repa_projector): Sequential(
+    9.97 M = 0.33% Params, 40.8 GMACs = 0.41% MACs, 698.57 us = 0.44% latency, 116.84 TFLOPS
+    (0): Linear(4.2 M = 0.14% Params, 17.18 GMACs = 0.17% MACs, 206.47 us = 0.13% latency, 166.41 TFLOPS, in_features=2048, out_features=2048, bias=True)
+    (1): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 34.57 us = 0.02% latency, 242.65 GFLOPS)
+    (2): Linear(4.2 M = 0.14% Params, 17.18 GMACs = 0.17% MACs, 171.9 us = 0.11% latency, 199.88 TFLOPS, in_features=2048, out_features=2048, bias=True)
+    (3): SiLU(0 = 0% Params, 0 MACs = 0% MACs, 36.24 us = 0.02% latency, 231.48 GFLOPS)
+    (4): Linear(1.57 M = 0.05% Params, 6.44 GMACs = 0.07% MACs, 148.06 us = 0.09% latency, 87.03 TFLOPS, in_features=2048, out_features=768, bias=True)
+  )
+)
+------------------------------------------------------------------------------

latest

@@ -0,0 +1 @@
+100000

zero_to_fp32.py

@@ -0,0 +1,604 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example: python zero_to_fp32.py . pytorch_model.bin
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage <= 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        state_dict = torch.load(file, map_location=device)
+
+        if BUFFER_NAMES not in state_dict:
+            raise ValueError(f"{file} is not a model state checkpoint")
+        buffer_names = state_dict[BUFFER_NAMES]
+        if debug:
+            print("Found buffers:", buffer_names)
+
+        # recover just the buffers while restoring them to fp32 if they were saved in fp16
+        buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+        param_shapes = state_dict[PARAM_SHAPES]
+
+        # collect parameters that are included in param_shapes
+        param_names = []
+        for s in param_shapes:
+            for name in s.keys():
+                param_names.append(name)
+
+        # update with frozen parameters
+        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+        if frozen_param_shapes is not None:
+            if debug:
+                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+            param_names += list(frozen_param_shapes.keys())
+
+        # handle shared params
+        shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+        ds_version = state_dict.get(DS_VERSION, None)
+
+        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+        z_model_state = zero_model_state(buffers=buffers,
+                                         param_shapes=param_shapes,
+                                         shared_params=shared_params,
+                                         ds_version=ds_version,
+                                         frozen_param_shapes=frozen_param_shapes,
+                                         frozen_param_fragments=frozen_param_fragments)
+        zero_model_states.append(z_model_state)
+
+    return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+
+    total_files = len(files)
+    state_dicts = []
+    for f in files:
+        state_dict = torch.load(f, map_location=device)
+        # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+        # and also handle the case where it was already removed by another helper script
+        state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+        state_dicts.append(state_dict)
+
+    if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage <= 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    if zero_stage <= 2:
+        fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+    elif zero_stage == 3:
+        # if there is more than one param group, there will be multiple flattened tensors - one
+        # flattened tensor per group - for simplicity merge them into a single tensor
+        #
+        # XXX: could make the script more memory efficient for when there are multiple groups - it
+        # will require matching the sub-lists of param_shapes for each param group flattened tensor
+
+        fp32_flat_groups = [
+            torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key], 0) for i in range(len(state_dicts))
+        ]
+
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_files = get_model_state_files(ds_checkpoint_dir)
+
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    if zero_stage <= 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _has_callable(obj, fn):
+    attr = getattr(obj, fn, None)
+    return callable(attr)
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum(
+        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in zero_model_states[0].frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+    avail_numel = fp32_flat_groups[0].numel() * world_size
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        avail_numel = fp32_flat_groups[0].numel() * world_size
+        print(f"Trainable params: Have {avail_numel} numels to process.")
+        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    for name, shape in param_shapes.items():
+
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # XXX: memory usage doubles here
+        state_dict[name] = torch.cat(
+            tuple(fp32_flat_groups[i].narrow(0, offset, partitioned_numel) for i in range(world_size)),
+            0).narrow(0, 0, unpartitioned_numel).view(shape)
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None, exclude_frozen_parameters=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+
+    Returns:
+        - pytorch ``state_dict``
+
+    Note: this approach may not work if your application doesn't have sufficient free CPU memory and
+    you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None, exclude_frozen_parameters=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+    """
+
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag, exclude_frozen_parameters)
+    print(f"Saving fp32 state dict to {output_file}")
+    torch.save(state_dict, output_file)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info(f"Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info(f"Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint_dir",
+                        type=str,
+                        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument(
+        "output_file",
+        type=str,
+        help="path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)")
+    parser.add_argument("-t",
+                        "--tag",
+                        type=str,
+                        default=None,
+                        help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+    parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
+                                               args.output_file,
+                                               tag=args.tag,
+                                               exclude_frozen_parameters=args.exclude_frozen_parameters)