Add files using upload-large-folder tool

LTA_openwebtext_dualt/logs/lta_lm1b_dirichlet_len1024_Cv_to_2v_gbs512_8gpu_20k_save1k_gumbelwatch_20260524.log

@@ -0,0 +1,687 @@
+
+*****************************************
+Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+*****************************************
+t-20260524091317-xb65t-worker-0:10324:10324 [0] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10324:10324 [0] NCCL INFO Bootstrap: Using eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10324:10324 [0] NCCL INFO cudaDriverVersion 12080
+t-20260524091317-xb65t-worker-0:10324:10324 [0] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260524091317-xb65t-worker-0:10324:10324 [0] NCCL INFO Comm config Blocking set to 1
+t-20260524091317-xb65t-worker-0:10325:10325 [1] NCCL INFO cudaDriverVersion 12080
+t-20260524091317-xb65t-worker-0:10325:10325 [1] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10325:10325 [1] NCCL INFO Bootstrap: Using eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10325:10325 [1] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260524091317-xb65t-worker-0:10325:10325 [1] NCCL INFO Comm config Blocking set to 1
+t-20260524091317-xb65t-worker-0:10329:10329 [5] NCCL INFO cudaDriverVersion 12080
+t-20260524091317-xb65t-worker-0:10329:10329 [5] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10329:10329 [5] NCCL INFO Bootstrap: Using eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10329:10329 [5] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260524091317-xb65t-worker-0:10329:10329 [5] NCCL INFO Comm config Blocking set to 1
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO P2P plugin v9 IBext_v9
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10327:10327 [3] NCCL INFO cudaDriverVersion 12080
+t-20260524091317-xb65t-worker-0:10327:10327 [3] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10327:10327 [3] NCCL INFO Bootstrap: Using eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10327:10327 [3] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260524091317-xb65t-worker-0:10327:10327 [3] NCCL INFO Comm config Blocking set to 1
+t-20260524091317-xb65t-worker-0:10328:10328 [4] NCCL INFO cudaDriverVersion 12080
+t-20260524091317-xb65t-worker-0:10328:10328 [4] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10328:10328 [4] NCCL INFO Bootstrap: Using eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10328:10328 [4] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260524091317-xb65t-worker-0:10328:10328 [4] NCCL INFO Comm config Blocking set to 1
+t-20260524091317-xb65t-worker-0:10330:10330 [6] NCCL INFO cudaDriverVersion 12080
+t-20260524091317-xb65t-worker-0:10330:10330 [6] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10330:10330 [6] NCCL INFO Bootstrap: Using eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10330:10330 [6] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260524091317-xb65t-worker-0:10330:10330 [6] NCCL INFO Comm config Blocking set to 1
+t-20260524091317-xb65t-worker-0:10326:10326 [2] NCCL INFO cudaDriverVersion 12080
+t-20260524091317-xb65t-worker-0:10326:10326 [2] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10326:10326 [2] NCCL INFO Bootstrap: Using eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10326:10326 [2] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260524091317-xb65t-worker-0:10326:10326 [2] NCCL INFO Comm config Blocking set to 1
+t-20260524091317-xb65t-worker-0:10331:10331 [7] NCCL INFO cudaDriverVersion 12080
+t-20260524091317-xb65t-worker-0:10331:10331 [7] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10331:10331 [7] NCCL INFO Bootstrap: Using eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10331:10331 [7] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260524091317-xb65t-worker-0:10331:10331 [7] NCCL INFO Comm config Blocking set to 1
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Using network IBext_v9
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO ncclCommInitRankConfig comm 0xd973700 rank 0 nranks 8 cudaDev 0 nvmlDev 0 busId 65040 commId 0xe94542814597d064 - Init START
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO P2P plugin v9 IBext_v9
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO P2P plugin v9 IBext_v9
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO P2P plugin v9 IBext_v9
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO P2P plugin v9 IBext_v9
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO P2P plugin v9 IBext_v9
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO P2P plugin v9 IBext_v9
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO P2P plugin v9 IBext_v9
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO Using network IBext_v9
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO Using network IBext_v9
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO Using network IBext_v9
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO Using network IBext_v9
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO Using network IBext_v9
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.96.52<0>
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO Using network IBext_v9
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO Using network IBext_v9
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO ncclCommInitRankConfig comm 0xb3a9dc0 rank 1 nranks 8 cudaDev 1 nvmlDev 1 busId 67020 commId 0xe94542814597d064 - Init START
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO ncclCommInitRankConfig comm 0xb803a00 rank 3 nranks 8 cudaDev 3 nvmlDev 3 busId 6b020 commId 0xe94542814597d064 - Init START
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO ncclCommInitRankConfig comm 0xa75cfc0 rank 5 nranks 8 cudaDev 5 nvmlDev 5 busId 71020 commId 0xe94542814597d064 - Init START
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO ncclCommInitRankConfig comm 0xa12b740 rank 4 nranks 8 cudaDev 4 nvmlDev 4 busId 6f020 commId 0xe94542814597d064 - Init START
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO ncclCommInitRankConfig comm 0xb514740 rank 6 nranks 8 cudaDev 6 nvmlDev 6 busId 73020 commId 0xe94542814597d064 - Init START
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO ncclCommInitRankConfig comm 0xb5394c0 rank 7 nranks 8 cudaDev 7 nvmlDev 7 busId 75020 commId 0xe94542814597d064 - Init START
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO ncclCommInitRankConfig comm 0x9d7a4c0 rank 2 nranks 8 cudaDev 2 nvmlDev 2 busId 69020 commId 0xe94542814597d064 - Init START
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO Bootstrap timings total 0.012864 (create 0.000026, send 0.000075, recv 0.007017, ring 0.005429, delay 0.000001)
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO Bootstrap timings total 0.008913 (create 0.000019, send 0.000061, recv 0.000120, ring 0.008428, delay 0.000001)
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO Bootstrap timings total 0.025766 (create 0.000020, send 0.000067, recv 0.016937, ring 0.001309, delay 0.000001)
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO Bootstrap timings total 0.001775 (create 0.000020, send 0.000069, recv 0.000095, ring 0.000093, delay 0.000001)
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO Bootstrap timings total 0.005908 (create 0.000020, send 0.000063, recv 0.003807, ring 0.001745, delay 0.000001)
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO Bootstrap timings total 0.105052 (create 0.000023, send 0.000062, recv 0.104563, ring 0.000086, delay 0.000001)
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO Bootstrap timings total 0.002241 (create 0.000023, send 0.000068, recv 0.000153, ring 0.001709, delay 0.000001)
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Bootstrap timings total 0.960418 (create 0.000026, send 0.000077, recv 0.855478, ring 0.001623, delay 0.000001)
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO MNNVL busId 0x6f020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO MNNVL busId 0x6b020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO MNNVL busId 0x71020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO MNNVL busId 0x73020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO MNNVL busId 0x69020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO MNNVL busId 0x67020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO MNNVL busId 0x75020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO MNNVL busId 0x65040 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO Setting affinity for GPU 2 to 03ffffff,ffffffff,ffffffff
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO NVLS multicast support is available on dev 2
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO Setting affinity for GPU 5 to 0fffff,ffffffff,ffffffff,fc000000,00000000,00000000
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO Setting affinity for GPU 7 to 0fffff,ffffffff,ffffffff,fc000000,00000000,00000000
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO Setting affinity for GPU 4 to 0fffff,ffffffff,ffffffff,fc000000,00000000,00000000
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO NVLS multicast support is available on dev 5
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO Setting affinity for GPU 3 to 03ffffff,ffffffff,ffffffff
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO NVLS multicast support is available on dev 4
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO NVLS multicast support is available on dev 3
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO Setting affinity for GPU 1 to 03ffffff,ffffffff,ffffffff
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Setting affinity for GPU 0 to 03ffffff,ffffffff,ffffffff
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO NVLS multicast support is available on dev 1
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO Setting affinity for GPU 6 to 0fffff,ffffffff,ffffffff,fc000000,00000000,00000000
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO NVLS multicast support is available on dev 0
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO NVLS multicast support is available on dev 6
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO NVLS multicast support is available on dev 7
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO comm 0xb514740 rank 6 nRanks 8 nNodes 1 localRanks 8 localRank 6 MNNVL 0
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO comm 0xa75cfc0 rank 5 nRanks 8 nNodes 1 localRanks 8 localRank 5 MNNVL 0
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO comm 0xa12b740 rank 4 nRanks 8 nNodes 1 localRanks 8 localRank 4 MNNVL 0
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO comm 0xb803a00 rank 3 nRanks 8 nNodes 1 localRanks 8 localRank 3 MNNVL 0
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO comm 0x9d7a4c0 rank 2 nRanks 8 nNodes 1 localRanks 8 localRank 2 MNNVL 0
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO comm 0xb3a9dc0 rank 1 nRanks 8 nNodes 1 localRanks 8 localRank 1 MNNVL 0
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO comm 0xd973700 rank 0 nRanks 8 nNodes 1 localRanks 8 localRank 0 MNNVL 0
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO comm 0xb5394c0 rank 7 nRanks 8 nNodes 1 localRanks 8 localRank 7 MNNVL 0
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO Trees [0] 5/-1/-1->4->3 [1] 5/-1/-1->4->3 [2] 5/-1/-1->4->3 [3] 5/-1/-1->4->3 [4] 5/-1/-1->4->3 [5] 5/-1/-1->4->3 [6] 5/-1/-1->4->3 [7] 5/-1/-1->4->3 [8] 5/-1/-1->4->3 [9] 5/-1/-1->4->3 [10] 5/-1/-1->4->3 [11] 5/-1/-1->4->3 [12] 5/-1/-1->4->3 [13] 5/-1/-1->4->3 [14] 5/-1/-1->4->3 [15] 5/-1/-1->4->3 [16] 5/-1/-1->4->3 [17] 5/-1/-1->4->3 [18] 5/-1/-1->4->3 [19] 5/-1/-1->4->3 [20] 5/-1/-1->4->3 [21] 5/-1/-1->4->3 [22] 5/-1/-1->4->3 [23] 5/-1/-1->4->3
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO P2P Chunksize set to 524288
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO Trees [0] 4/-1/-1->3->2 [1] 4/-1/-1->3->2 [2] 4/-1/-1->3->2 [3] 4/-1/-1->3->2 [4] 4/-1/-1->3->2 [5] 4/-1/-1->3->2 [6] 4/-1/-1->3->2 [7] 4/-1/-1->3->2 [8] 4/-1/-1->3->2 [9] 4/-1/-1->3->2 [10] 4/-1/-1->3->2 [11] 4/-1/-1->3->2 [12] 4/-1/-1->3->2 [13] 4/-1/-1->3->2 [14] 4/-1/-1->3->2 [15] 4/-1/-1->3->2 [16] 4/-1/-1->3->2 [17] 4/-1/-1->3->2 [18] 4/-1/-1->3->2 [19] 4/-1/-1->3->2 [20] 4/-1/-1->3->2 [21] 4/-1/-1->3->2 [22] 4/-1/-1->3->2 [23] 4/-1/-1->3->2
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 00/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO Trees [0] 6/-1/-1->5->4 [1] 6/-1/-1->5->4 [2] 6/-1/-1->5->4 [3] 6/-1/-1->5->4 [4] 6/-1/-1->5->4 [5] 6/-1/-1->5->4 [6] 6/-1/-1->5->4 [7] 6/-1/-1->5->4 [8] 6/-1/-1->5->4 [9] 6/-1/-1->5->4 [10] 6/-1/-1->5->4 [11] 6/-1/-1->5->4 [12] 6/-1/-1->5->4 [13] 6/-1/-1->5->4 [14] 6/-1/-1->5->4 [15] 6/-1/-1->5->4 [16] 6/-1/-1->5->4 [17] 6/-1/-1->5->4 [18] 6/-1/-1->5->4 [19] 6/-1/-1->5->4 [20] 6/-1/-1->5->4 [21] 6/-1/-1->5->4 [22] 6/-1/-1->5->4 [23] 6/-1/-1->5->4
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO P2P Chunksize set to 524288
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO Trees [0] 7/-1/-1->6->5 [1] 7/-1/-1->6->5 [2] 7/-1/-1->6->5 [3] 7/-1/-1->6->5 [4] 7/-1/-1->6->5 [5] 7/-1/-1->6->5 [6] 7/-1/-1->6->5 [7] 7/-1/-1->6->5 [8] 7/-1/-1->6->5 [9] 7/-1/-1->6->5 [10] 7/-1/-1->6->5 [11] 7/-1/-1->6->5 [12] 7/-1/-1->6->5 [13] 7/-1/-1->6->5 [14] 7/-1/-1->6->5 [15] 7/-1/-1->6->5 [16] 7/-1/-1->6->5 [17] 7/-1/-1->6->5 [18] 7/-1/-1->6->5 [19] 7/-1/-1->6->5 [20] 7/-1/-1->6->5 [21] 7/-1/-1->6->5 [22] 7/-1/-1->6->5 [23] 7/-1/-1->6->5
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 01/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO P2P Chunksize set to 524288
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 02/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO P2P Chunksize set to 524288
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 03/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 04/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 05/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 06/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 07/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO Trees [0] 2/-1/-1->1->0 [1] 2/-1/-1->1->0 [2] 2/-1/-1->1->0 [3] 2/-1/-1->1->0 [4] 2/-1/-1->1->0 [5] 2/-1/-1->1->0 [6] 2/-1/-1->1->0 [7] 2/-1/-1->1->0 [8] 2/-1/-1->1->0 [9] 2/-1/-1->1->0 [10] 2/-1/-1->1->0 [11] 2/-1/-1->1->0 [12] 2/-1/-1->1->0 [13] 2/-1/-1->1->0 [14] 2/-1/-1->1->0 [15] 2/-1/-1->1->0 [16] 2/-1/-1->1->0 [17] 2/-1/-1->1->0 [18] 2/-1/-1->1->0 [19] 2/-1/-1->1->0 [20] 2/-1/-1->1->0 [21] 2/-1/-1->1->0 [22] 2/-1/-1->1->0 [23] 2/-1/-1->1->0
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO Trees [0] 3/-1/-1->2->1 [1] 3/-1/-1->2->1 [2] 3/-1/-1->2->1 [3] 3/-1/-1->2->1 [4] 3/-1/-1->2->1 [5] 3/-1/-1->2->1 [6] 3/-1/-1->2->1 [7] 3/-1/-1->2->1 [8] 3/-1/-1->2->1 [9] 3/-1/-1->2->1 [10] 3/-1/-1->2->1 [11] 3/-1/-1->2->1 [12] 3/-1/-1->2->1 [13] 3/-1/-1->2->1 [14] 3/-1/-1->2->1 [15] 3/-1/-1->2->1 [16] 3/-1/-1->2->1 [17] 3/-1/-1->2->1 [18] 3/-1/-1->2->1 [19] 3/-1/-1->2->1 [20] 3/-1/-1->2->1 [21] 3/-1/-1->2->1 [22] 3/-1/-1->2->1 [23] 3/-1/-1->2->1
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO Trees [0] -1/-1/-1->7->6 [1] -1/-1/-1->7->6 [2] -1/-1/-1->7->6 [3] -1/-1/-1->7->6 [4] -1/-1/-1->7->6 [5] -1/-1/-1->7->6 [6] -1/-1/-1->7->6 [7] -1/-1/-1->7->6 [8] -1/-1/-1->7->6 [9] -1/-1/-1->7->6 [10] -1/-1/-1->7->6 [11] -1/-1/-1->7->6 [12] -1/-1/-1->7->6 [13] -1/-1/-1->7->6 [14] -1/-1/-1->7->6 [15] -1/-1/-1->7->6 [16] -1/-1/-1->7->6 [17] -1/-1/-1->7->6 [18] -1/-1/-1->7->6 [19] -1/-1/-1->7->6 [20] -1/-1/-1->7->6 [21] -1/-1/-1->7->6 [22] -1/-1/-1->7->6 [23] -1/-1/-1->7->6
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 08/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO P2P Chunksize set to 524288
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO P2P Chunksize set to 524288
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 09/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO P2P Chunksize set to 524288
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 10/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 11/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 12/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 13/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 14/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 15/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 16/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 17/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 18/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 19/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 20/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 21/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 22/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Channel 23/24 : 0 1 2 3 4 5 6 7
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Trees [0] 1/-1/-1->0->-1 [1] 1/-1/-1->0->-1 [2] 1/-1/-1->0->-1 [3] 1/-1/-1->0->-1 [4] 1/-1/-1->0->-1 [5] 1/-1/-1->0->-1 [6] 1/-1/-1->0->-1 [7] 1/-1/-1->0->-1 [8] 1/-1/-1->0->-1 [9] 1/-1/-1->0->-1 [10] 1/-1/-1->0->-1 [11] 1/-1/-1->0->-1 [12] 1/-1/-1->0->-1 [13] 1/-1/-1->0->-1 [14] 1/-1/-1->0->-1 [15] 1/-1/-1->0->-1 [16] 1/-1/-1->0->-1 [17] 1/-1/-1->0->-1 [18] 1/-1/-1->0->-1 [19] 1/-1/-1->0->-1 [20] 1/-1/-1->0->-1 [21] 1/-1/-1->0->-1 [22] 1/-1/-1->0->-1 [23] 1/-1/-1->0->-1
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO P2P Chunksize set to 524288
+t-20260524091317-xb65t-worker-0:10329:10476 [5] NCCL INFO [Proxy Service] Device 5 CPU core 94
+t-20260524091317-xb65t-worker-0:10329:10477 [5] NCCL INFO [Proxy Service UDS] Device 5 CPU core 96
+t-20260524091317-xb65t-worker-0:10327:10478 [3] NCCL INFO [Proxy Service] Device 3 CPU core 14
+t-20260524091317-xb65t-worker-0:10327:10479 [3] NCCL INFO [Proxy Service UDS] Device 3 CPU core 18
+t-20260524091317-xb65t-worker-0:10328:10480 [4] NCCL INFO [Proxy Service] Device 4 CPU core 168
+t-20260524091317-xb65t-worker-0:10330:10482 [6] NCCL INFO [Proxy Service] Device 6 CPU core 96
+t-20260524091317-xb65t-worker-0:10328:10481 [4] NCCL INFO [Proxy Service UDS] Device 4 CPU core 172
+t-20260524091317-xb65t-worker-0:10330:10483 [6] NCCL INFO [Proxy Service UDS] Device 6 CPU core 94
+t-20260524091317-xb65t-worker-0:10325:10484 [1] NCCL INFO [Proxy Service] Device 1 CPU core 2
+t-20260524091317-xb65t-worker-0:10325:10485 [1] NCCL INFO [Proxy Service UDS] Device 1 CPU core 4
+t-20260524091317-xb65t-worker-0:10331:10486 [7] NCCL INFO [Proxy Service] Device 7 CPU core 94
+t-20260524091317-xb65t-worker-0:10331:10487 [7] NCCL INFO [Proxy Service UDS] Device 7 CPU core 96
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Check P2P Type intraNodeP2pSupport 1 directMode 0
+t-20260524091317-xb65t-worker-0:10324:10488 [0] NCCL INFO [Proxy Service] Device 0 CPU core 22
+t-20260524091317-xb65t-worker-0:10326:10489 [2] NCCL INFO [Proxy Service] Device 2 CPU core 2
+t-20260524091317-xb65t-worker-0:10324:10490 [0] NCCL INFO [Proxy Service UDS] Device 0 CPU core 18
+t-20260524091317-xb65t-worker-0:10326:10491 [2] NCCL INFO [Proxy Service UDS] Device 2 CPU core 4
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO CC Off, workFifoBytes 1048576
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO ncclCommInitRankConfig comm 0xb803a00 rank 3 nranks 8 cudaDev 3 nvmlDev 3 busId 6b020 commId 0xe94542814597d064 - Init COMPLETE
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO ncclCommInitRankConfig comm 0xa12b740 rank 4 nranks 8 cudaDev 4 nvmlDev 4 busId 6f020 commId 0xe94542814597d064 - Init COMPLETE
+t-20260524091317-xb65t-worker-0:10327:10400 [3] NCCL INFO Init timings - ncclCommInitRankConfig: rank 3 nranks 8 total 2.14 (kernels 0.22, alloc 0.90, bootstrap 0.03, allgathers 0.01, topo 0.54, graphs 0.01, connections 0.41, rest 0.02)
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO ncclCommInitRankConfig comm 0xb5394c0 rank 7 nranks 8 cudaDev 7 nvmlDev 7 busId 75020 commId 0xe94542814597d064 - Init COMPLETE
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO ncclCommInitRankConfig comm 0xa75cfc0 rank 5 nranks 8 cudaDev 5 nvmlDev 5 busId 71020 commId 0xe94542814597d064 - Init COMPLETE
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO ncclCommInitRankConfig comm 0x9d7a4c0 rank 2 nranks 8 cudaDev 2 nvmlDev 2 busId 69020 commId 0xe94542814597d064 - Init COMPLETE
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO ncclCommInitRankConfig comm 0xb3a9dc0 rank 1 nranks 8 cudaDev 1 nvmlDev 1 busId 67020 commId 0xe94542814597d064 - Init COMPLETE
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO ncclCommInitRankConfig comm 0xd973700 rank 0 nranks 8 cudaDev 0 nvmlDev 0 busId 65040 commId 0xe94542814597d064 - Init COMPLETE
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO ncclCommInitRankConfig comm 0xb514740 rank 6 nranks 8 cudaDev 6 nvmlDev 6 busId 73020 commId 0xe94542814597d064 - Init COMPLETE
+t-20260524091317-xb65t-worker-0:10328:10401 [4] NCCL INFO Init timings - ncclCommInitRankConfig: rank 4 nranks 8 total 2.13 (kernels 0.22, alloc 0.91, bootstrap 0.01, allgathers 0.01, topo 0.54, graphs 0.01, connections 0.41, rest 0.03)
+t-20260524091317-xb65t-worker-0:10331:10408 [7] NCCL INFO Init timings - ncclCommInitRankConfig: rank 7 nranks 8 total 2.11 (kernels 0.22, alloc 0.90, bootstrap 0.00, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.41, rest 0.03)
+t-20260524091317-xb65t-worker-0:10329:10398 [5] NCCL INFO Init timings - ncclCommInitRankConfig: rank 5 nranks 8 total 2.15 (kernels 0.23, alloc 0.91, bootstrap 0.01, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.42, rest 0.02)
+t-20260524091317-xb65t-worker-0:10324:10396 [0] NCCL INFO Init timings - ncclCommInitRankConfig: rank 0 nranks 8 total 2.36 (kernels 0.21, alloc 0.20, bootstrap 0.96, allgathers 0.01, topo 0.54, graphs 0.01, connections 0.41, rest 0.03)
+t-20260524091317-xb65t-worker-0:10325:10397 [1] NCCL INFO Init timings - ncclCommInitRankConfig: rank 1 nranks 8 total 2.20 (kernels 0.27, alloc 0.84, bootstrap 0.11, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.41, rest 0.03)
+t-20260524091317-xb65t-worker-0:10326:10406 [2] NCCL INFO Init timings - ncclCommInitRankConfig: rank 2 nranks 8 total 2.12 (kernels 0.22, alloc 0.91, bootstrap 0.00, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.41, rest 0.03)
+t-20260524091317-xb65t-worker-0:10330:10404 [6] NCCL INFO Init timings - ncclCommInitRankConfig: rank 6 nranks 8 total 2.13 (kernels 0.22, alloc 0.91, bootstrap 0.01, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.41, rest 0.03)
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 00/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 00/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 00/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 01/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 00/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 01/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 01/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 02/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 01/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 02/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 02/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 03/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 02/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 03/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 04/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 05/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 00/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 06/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 01/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 07/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 02/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 08/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 03/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 09/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 04/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 10/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 05/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 11/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 06/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 12/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 07/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 03/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 13/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 03/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 04/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 00/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 08/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 04/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 14/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 04/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 05/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 01/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 09/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 05/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 15/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 05/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 06/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 02/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 10/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 06/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 16/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 06/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 07/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 03/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 00/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 04/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 11/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 00/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 01/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 05/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 12/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 01/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 02/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 06/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 13/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 02/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 07/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 03/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 07/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 14/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 03/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 08/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 04/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 08/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 15/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 04/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 09/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 05/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 09/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 16/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 05/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 10/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 06/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 10/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 17/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 17/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 07/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 06/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 08/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 11/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 07/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 11/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 18/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 18/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 08/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 07/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 09/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 12/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 08/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 12/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 19/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 19/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 09/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 08/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 10/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 13/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 09/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 13/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 20/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 20/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 10/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 09/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 11/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 14/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 10/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 14/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 21/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 21/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 11/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 10/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 12/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 15/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 11/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 15/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 22/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 22/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 12/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 11/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 13/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 16/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 12/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 16/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Channel 23/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Channel 23/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 13/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 12/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 14/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 17/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 13/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 17/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 14/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 13/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 15/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 18/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 14/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 18/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 15/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 19/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 14/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 16/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 15/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 19/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 16/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 20/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 15/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 17/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 16/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 20/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 17/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 21/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 16/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 18/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 17/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 21/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 18/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 22/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 17/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 19/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 18/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 22/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 19/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Channel 23/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 18/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 20/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 19/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Channel 23/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 19/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 20/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 21/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 20/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 21/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 22/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 20/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 22/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 21/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Channel 23/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 21/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Channel 23/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 22/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 22/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Channel 23/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Channel 23/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260524091317-xb65t-worker-0:10330:10497 [6] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260524091317-xb65t-worker-0:10331:10499 [7] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260524091317-xb65t-worker-0:10329:10492 [5] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260524091317-xb65t-worker-0:10325:10498 [1] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260524091317-xb65t-worker-0:10326:10495 [2] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260524091317-xb65t-worker-0:10327:10493 [3] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260524091317-xb65t-worker-0:10324:10494 [0] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260524091317-xb65t-worker-0:10328:10496 [4] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+{
+  "device": "cuda:0",
+  "rank": 0,
+  "world_size": 8,
+  "samples": "wrapped_stream",
+  "vocab_size": 30522,
+  "tokenizer_vocab_size": 30522,
+  "save_dir": "runs/lta_lm1b_dirichlet_len1024_Cv_to_2v_gbs512_8gpu_20k_save1k_gumbelwatch_20260524",
+  "max_len": 1024,
+  "effective_model_max_len": 1024,
+  "batch_size": 2,
+  "grad_accum": 32,
+  "effective_batch_size": 512,
+  "global_batch_size": 512,
+  "lr_schedule": "constant_warmup",
+  "optimizer": "adamw",
+  "epochs": 0.0,
+  "steps_per_epoch": 0,
+  "total_steps": 20000,
+  "warmup_steps": 2500,
+  "warmup_epochs": -1.0,
+  "min_lr": 6e-05,
+  "weight_decay": 0.0,
+  "output_weight_decay": -1.0,
+  "adamw_param_groups": "nanogpt",
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.999,
+  "adam_eps": 1e-08,
+  "muon_impl": "legacy",
+  "muon_momentum": 0.95,
+  "muon_ns_steps": 5,
+  "muon_update_scale": 1.0,
+  "muon_nesterov": false,
+  "muon_width_scale": false,
+  "muon_grouping": "",
+  "muon_param_count": 0,
+  "muon_adam_param_count": 0,
+  "muon_param_names": [],
+  "muon_adam_param_names": [],
+  "muon_effective_nesterov": false,
+  "muon_effective_width_scale": false,
+  "muon_effective_weight_decay": 0.0,
+  "muon_adam_fallback_nesterov": false,
+  "muon_adam_fallback_weight_decay": 0.0,
+  "ema_decay": 0.0,
+  "ema_start_step": 0,
+  "model_type": "ddit",
+  "ddit_mlp_type": "gelu",
+  "block_anchor_every": 0,
+  "block_anchor_init_std": 0.02,
+  "bos_anchor_every": 0,
+  "bos_anchor_token_id": -1,
+  "bos_anchor_extra_len": 0,
+  "abs_pos_embed": false,
+  "abs_pos_init_std": 0.02,
+  "elf_num_time_tokens": 4,
+  "elf_num_model_mode_tokens": 0,
+  "qk_norm": true,
+  "output_bias": false,
+  "output_init_std": -1.0,
+  "norm_type": "rmsnorm",
+  "target_loss": "hard_ce",
+  "linear_soft_target_power": 1.0,
+  "linear_soft_target_min_conf": 0.0,
+  "linear_soft_target_max_conf": 1.0,
+  "t_sampling_mode": "uniform",
+  "t_sampling_power": 1.0,
+  "t_sampling_eps": 0.0001,
+  "t_sampling_logit_mean": -1.5,
+  "t_sampling_logit_std": 0.8,
+  "t_sampling_gumbel_loc": 2.2,
+  "t_sampling_gumbel_scale": 0.8,
+  "dual_t": true,
+  "corrupt_t_mode": "independent",
+  "corrupt_min_t": 0.0,
+  "corrupt_max_t": 1.0,
+  "prefix_block_prob": 0.0,
+  "prefix_block_len": 128,
+  "block_ar_two_stream": false,
+  "block_ar_block_len": 128,
+  "mask_ratio_floor_schedule": "none",
+  "dirichlet_endpoint_mode": "categorical_dual_t",
+  "dirichlet_semantic_t_mode": "same",
+  "dirichlet_semantic_t_value": 0.0,
+  "dirichlet_semantic_t_curve": "linear",
+  "dirichlet_semantic_t_power": 1.0,
+  "dirichlet_support_t_curve": "linear",
+  "dirichlet_support_t_power": 1.0,
+  "endpoint_sequence_random_prob_alpha": 0.0,
+  "categorical_wrong_from_full_vocab": true,
+  "categorical_wrong_from_batch_valid_tokens": false,
+  "categorical_wrong_basin_token_ids": "",
+  "categorical_wrong_basin_prob": 0.0,
+  "categorical_wrong_unigram_prob": 0.0,
+  "categorical_wrong_uniform_prob": 0.0,
+  "categorical_wrong_prob_floor": 0.0,
+  "categorical_gold_prob_floor": 0.0,
+  "categorical_gold_prob_ceil": 1.0,
+  "categorical_wrong_corpus_unigram_path": "",
+  "categorical_wrong_corpus_unigram_alpha": 1.0,
+  "categorical_wrong_basin_shared_prob": 0.0,
+  "categorical_wrong_unigram_shared_prob": 0.0,
+  "mask_mixture_original_prob": 0.0,
+  "mask_mixture_lowk_prob": 0.0,
+  "mask_mixture_lowcorrupt_prob": 0.0,
+  "mask_mixture_block_prob": 0.0,
+  "mask_mixture_all_prob": 0.0,
+  "mask_mixture_lowk_clean_tokens": "1,2,4,8,16,32,64",
+  "mask_mixture_lowcorrupt_tokens": "1,2,4,8,16,32,64",
+  "mask_mixture_block_tokens": "64,128",
+  "simplex_bridge_sampler": "dirichlet",
+  "logistic_normal_sigma_min": 0.18,
+  "logistic_normal_sigma_max": 2.2,
+  "logistic_normal_tau_min": 0.65,
+  "logistic_normal_tau_max": 1.15,
+  "torch_compile": false,
+  "compile_mode": "max-autotune",
+  "state_format": "prob",
+  "meanflow_weight": 0.0,
+  "rollout_train_prob": 0.0,
+  "rollout_train_steps": 1,
+  "rollout_train_steps_min": -1,
+  "rollout_train_infer_steps": 64,
+  "rollout_train_time_mode": "fixed_steps",
+  "rollout_train_s_dist": "uniform",
+  "rollout_train_s_min_frac": 0.0,
+  "rollout_train_s_max_frac": 0.125,
+  "rollout_train_s_beta_alpha": 2.0,
+  "rollout_train_s_beta_beta": 6.0,
+  "rollout_train_temp": 1.0,
+  "rollout_train_max_gamma": 1.0,
+  "rollout_train_rule": "flowmap",
+  "rollout_train_corrupt_only": true,
+  "rollout_train_samplewise": false,
+  "rollout_train_compute_always": false,
+  "rollout_train_keep_grad": false,
+  "rollout_train_sync_t": false,
+  "rollout_train_state_mix_mode": "final",
+  "rollout_train_state_mix_alpha": 0.5,
+  "bridge_noise_init": "logistic_normal",
+  "noise_sigma": -1.0,
+  "allow_tf32": true,
+  "activation_checkpointing": false,
+  "activation_checkpoint_interval": 1,
+  "activation_checkpoint_scope": "block",
+  "ddp_static_graph": false,
+  "ddp_gradient_as_bucket_view": true,
+  "blocking_data_transfer": false,
+  "dataloader_prefetch_factor": 2,
+  "full_train_stats": false,
+  "tokenized_hf": false,
+  "tokenized_pad_token": "pad",
+  "elf_conditional_hf": false,
+  "record_pad_truncate": false,
+  "record_add_eos": false,
+  "record_add_special_tokens": false,
+  "record_pad_token": "pad",
+  "record_shuffle_buffer": 10000,
+  "wrap": true,
+  "wrap_mode": "stream",
+  "wrap_record_buffer_size": 200,
+  "owt_cached_chunks": false,
+  "owt_chunk_cache_dir": "",
+  "owt_chunk_cache_rebuild": false,
+  "owt_chunk_cache_write_batch": 4096,
+  "owt_exact_repeat_per_chunk": 0,
+  "online_chunk_shuffle": false,
+  "online_chunk_shuffle_buffer": 10000,
+  "openwebtext_split": "all",
+  "detokenizer": "auto",
+  "resolved_detokenizer": "lm1b",
+  "num_workers": 0,
+  "latest_every": 1000,
+  "resume_path": ""
+}
+t-20260524091317-xb65t-worker-0:10324:10604 [0] NCCL INFO NVLS comm 0xd973700 headRank 0 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260524091317-xb65t-worker-0:10325:10605 [1] NCCL INFO NVLS comm 0xb3a9dc0 headRank 1 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260524091317-xb65t-worker-0:10326:10606 [2] NCCL INFO NVLS comm 0x9d7a4c0 headRank 2 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260524091317-xb65t-worker-0:10328:10607 [4] NCCL INFO NVLS comm 0xa12b740 headRank 4 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260524091317-xb65t-worker-0:10327:10608 [3] NCCL INFO NVLS comm 0xb803a00 headRank 3 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260524091317-xb65t-worker-0:10329:10609 [5] NCCL INFO NVLS comm 0xa75cfc0 headRank 5 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260524091317-xb65t-worker-0:10330:10610 [6] NCCL INFO NVLS comm 0xb514740 headRank 6 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260524091317-xb65t-worker-0:10331:10611 [7] NCCL INFO NVLS comm 0xb5394c0 headRank 7 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+step=100 micro_steps=3200 elapsed=170.6s lr=1.212000e-05 loss=10.1456 loss_recon=10.1456 loss_meanflow=0.0000 mean_model_t=0.4981 mean_corrupt_t=0.4994 mean_loss_t_weight=1.0000 linear_soft_target_mean_conf=0.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.0000 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 acc_all=0.5713 corrupt_frac=0.5504 acc_corrupt=0.3807 loss_corrupt=10.1456 wrong_frac=0.5012 init_acc_corrupt=0.4988 acc_corrupt_t_0p0_0p2=0.0791 corrupt_frac_t_0p0_0p2=0.5561 acc_corrupt_t_0p6_0p8=0.5293 corrupt_frac_t_0p6_0p8=0.5521 out_w_norm=1.0058 out_g_norm=1.0429 acc_corrupt_t_0p2_0p4=0.2091 corrupt_frac_t_0p2_0p4=0.5638 acc_corrupt_t_0p8_1p0=0.7259 corrupt_frac_t_0p8_1p0=0.5535 acc_corrupt_t_0p4_0p6=0.3628 corrupt_frac_t_0p4_0p6=0.5436 loss_all=9.6693 init_gold_top10=0.5832 init_gold_top100=0.5832
+step=200 micro_steps=6400 elapsed=165.6s lr=2.412000e-05 loss=8.8510 loss_recon=8.8510 loss_meanflow=0.0000 mean_model_t=0.5027 mean_corrupt_t=0.4912 mean_loss_t_weight=1.0000 linear_soft_target_mean_conf=0.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.0000 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 acc_all=0.1270 corrupt_frac=0.5485 acc_corrupt=0.0923 loss_corrupt=8.8510 wrong_frac=0.5090 init_acc_corrupt=0.4910 acc_corrupt_t_0p6_0p8=0.1149 corrupt_frac_t_0p6_0p8=0.5704 acc_corrupt_t_0p8_1p0=0.1501 corrupt_frac_t_0p8_1p0=0.5525 out_w_norm=6.9327 out_g_norm=1.5119 acc_corrupt_t_0p0_0p2=0.0460 corrupt_frac_t_0p0_0p2=0.5588 acc_corrupt_t_0p4_0p6=0.0890 corrupt_frac_t_0p4_0p6=0.5557 acc_corrupt_t_0p2_0p4=0.0667 corrupt_frac_t_0p2_0p4=0.5525 loss_all=7.9347 init_gold_top10=0.3662 init_gold_top100=0.3662
+step=300 micro_steps=9600 elapsed=169.5s lr=3.612000e-05 loss=6.9617 loss_recon=6.9617 loss_meanflow=0.0000 mean_model_t=0.5020 mean_corrupt_t=0.4988 mean_loss_t_weight=1.0000 linear_soft_target_mean_conf=0.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.0000 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 acc_all=0.1500 corrupt_frac=0.5494 acc_corrupt=0.1156 loss_corrupt=6.9617 wrong_frac=0.5005 init_acc_corrupt=0.4996 acc_corrupt_t_0p6_0p8=0.1446 corrupt_frac_t_0p6_0p8=0.5593 out_w_norm=13.3103 out_g_norm=1.1168 acc_corrupt_t_0p2_0p4=0.0853 corrupt_frac_t_0p2_0p4=0.5636 acc_corrupt_t_0p8_1p0=0.1791 corrupt_frac_t_0p8_1p0=0.5610 acc_corrupt_t_0p4_0p6=0.1135 corrupt_frac_t_0p4_0p6=0.5426 acc_corrupt_t_0p0_0p2=0.0557 corrupt_frac_t_0p0_0p2=0.5508 loss_all=5.0093 init_gold_top10=0.5300 init_gold_top100=0.5309
+step=400 micro_steps=12800 elapsed=173.6s lr=4.812000e-05 loss=4.3888 loss_recon=4.3888 loss_meanflow=0.0000 mean_model_t=0.4947 mean_corrupt_t=0.4990 mean_loss_t_weight=1.0000 linear_soft_target_mean_conf=0.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.0000 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 acc_all=0.6375 corrupt_frac=0.5480 acc_corrupt=0.4451 loss_corrupt=4.3888 wrong_frac=0.5022 init_acc_corrupt=0.4978 acc_corrupt_t_0p0_0p2=0.1092 corrupt_frac_t_0p0_0p2=0.5475 out_w_norm=18.9742 out_g_norm=0.4668 acc_corrupt_t_0p6_0p8=0.6186 corrupt_frac_t_0p6_0p8=0.5556 acc_corrupt_t_0p8_1p0=0.7924 corrupt_frac_t_0p8_1p0=0.5553 acc_corrupt_t_0p4_0p6=0.4457 corrupt_frac_t_0p4_0p6=0.5529 acc_corrupt_t_0p2_0p4=0.2734 corrupt_frac_t_0p2_0p4=0.5768 loss_all=3.0366 init_gold_top10=0.2575 init_gold_top100=0.2605
+step=500 micro_steps=16000 elapsed=172.9s lr=6.012000e-05 loss=3.8340 loss_recon=3.8340 loss_meanflow=0.0000 mean_model_t=0.4983 mean_corrupt_t=0.5041 mean_loss_t_weight=1.0000 linear_soft_target_mean_conf=0.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.0000 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 acc_all=0.7222 corrupt_frac=0.5470 acc_corrupt=0.5047 loss_corrupt=3.8340 wrong_frac=0.4975 init_acc_corrupt=0.5025 acc_corrupt_t_0p4_0p6=0.4986 corrupt_frac_t_0p4_0p6=0.5556 acc_corrupt_t_0p8_1p0=0.8904 corrupt_frac_t_0p8_1p0=0.5599 out_w_norm=22.4888 out_g_norm=0.3746 acc_corrupt_t_0p2_0p4=0.3076 corrupt_frac_t_0p2_0p4=0.5515 acc_corrupt_t_0p6_0p8=0.6927 corrupt_frac_t_0p6_0p8=0.5543 acc_corrupt_t_0p0_0p2=0.1224 corrupt_frac_t_0p0_0p2=0.5550 loss_all=2.7828 init_gold_top10=0.2262 init_gold_top100=0.2303
+step=600 micro_steps=19200 elapsed=167.3s lr=7.212000e-05 loss=3.7715 loss_recon=3.7715 loss_meanflow=0.0000 mean_model_t=0.5091 mean_corrupt_t=0.5016 mean_loss_t_weight=1.0000 linear_soft_target_mean_conf=0.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.0000 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 acc_all=0.7281 corrupt_frac=0.5560 acc_corrupt=0.5141 loss_corrupt=3.7715 wrong_frac=0.4967 init_acc_corrupt=0.5033 acc_corrupt_t_0p2_0p4=0.3151 corrupt_frac_t_0p2_0p4=0.5598 acc_corrupt_t_0p8_1p0=0.8966 corrupt_frac_t_0p8_1p0=0.5597 out_w_norm=24.2086 out_g_norm=0.4152 acc_corrupt_t_0p0_0p2=0.1330 corrupt_frac_t_0p0_0p2=0.5436 acc_corrupt_t_0p6_0p8=0.7071 corrupt_frac_t_0p6_0p8=0.5590 acc_corrupt_t_0p4_0p6=0.5077 corrupt_frac_t_0p4_0p6=0.5540 loss_all=2.5931 init_gold_top10=0.3047 init_gold_top100=0.3067
+step=700 micro_steps=22400 elapsed=171.0s lr=8.412000e-05 loss=3.6355 loss_recon=3.6355 loss_meanflow=0.0000 mean_model_t=0.4975 mean_corrupt_t=0.5052 mean_loss_t_weight=1.0000 linear_soft_target_mean_conf=0.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.0000 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 acc_all=0.7413 corrupt_frac=0.5475 acc_corrupt=0.5296 loss_corrupt=3.6355 wrong_frac=0.4965 init_acc_corrupt=0.5035 acc_corrupt_t_0p8_1p0=0.9016 corrupt_frac_t_0p8_1p0=0.5505 out_w_norm=25.3922 out_g_norm=0.4602 acc_corrupt_t_0p2_0p4=0.3379 corrupt_frac_t_0p2_0p4=0.5590 acc_corrupt_t_0p4_0p6=0.5255 corrupt_frac_t_0p4_0p6=0.5618 acc_corrupt_t_0p0_0p2=0.1564 corrupt_frac_t_0p0_0p2=0.5508 acc_corrupt_t_0p6_0p8=0.7125 corrupt_frac_t_0p6_0p8=0.5485 loss_all=0.9812 init_gold_top10=0.7266 init_gold_top100=0.7276
+step=800 micro_steps=25600 elapsed=173.0s lr=9.612000e-05 loss=3.4932 loss_recon=3.4932 loss_meanflow=0.0000 mean_model_t=0.4942 mean_corrupt_t=0.4965 mean_loss_t_weight=1.0000 linear_soft_target_mean_conf=0.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.0000 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 acc_all=0.7437 corrupt_frac=0.5486 acc_corrupt=0.5348 loss_corrupt=3.4932 wrong_frac=0.5036 init_acc_corrupt=0.4964 acc_corrupt_t_0p0_0p2=0.1709 corrupt_frac_t_0p0_0p2=0.5720 acc_corrupt_t_0p8_1p0=0.9047 corrupt_frac_t_0p8_1p0=0.5536 out_w_norm=26.5098 out_g_norm=0.5269 acc_corrupt_t_0p4_0p6=0.5399 corrupt_frac_t_0p4_0p6=0.5477 acc_corrupt_t_0p2_0p4=0.3559 corrupt_frac_t_0p2_0p4=0.5440 acc_corrupt_t_0p6_0p8=0.7240 corrupt_frac_t_0p6_0p8=0.5552 loss_all=0.2252 init_gold_top10=0.9130 init_gold_top100=0.9130
+step=900 micro_steps=28800 elapsed=170.7s lr=1.081200e-04 loss=3.3160 loss_recon=3.3160 loss_meanflow=0.0000 mean_model_t=0.5031 mean_corrupt_t=0.5038 mean_loss_t_weight=1.0000 linear_soft_target_mean_conf=0.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.0000 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 acc_all=0.7492 corrupt_frac=0.5523 acc_corrupt=0.5479 loss_corrupt=3.3160 wrong_frac=0.4972 init_acc_corrupt=0.5028 acc_corrupt_t_0p4_0p6=0.5454 corrupt_frac_t_0p4_0p6=0.5577 out_w_norm=27.5695 out_g_norm=0.5297 acc_corrupt_t_0p2_0p4=0.3649 corrupt_frac_t_0p2_0p4=0.5480 acc_corrupt_t_0p6_0p8=0.7308 corrupt_frac_t_0p6_0p8=0.5478 acc_corrupt_t_0p8_1p0=0.9107 corrupt_frac_t_0p8_1p0=0.5578 acc_corrupt_t_0p0_0p2=0.1776 corrupt_frac_t_0p0_0p2=0.5622 loss_all=2.2426 init_gold_top10=0.6347 init_gold_top100=0.6358
+step=1000 micro_steps=32000 elapsed=168.6s lr=1.201200e-04 loss=3.2368 loss_recon=3.2368 loss_meanflow=0.0000 mean_model_t=0.5029 mean_corrupt_t=0.4996 mean_loss_t_weight=1.0000 linear_soft_target_mean_conf=0.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.0000 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 acc_all=0.7512 corrupt_frac=0.5516 acc_corrupt=0.5509 loss_corrupt=3.2368 wrong_frac=0.5002 init_acc_corrupt=0.4999 acc_corrupt_t_0p6_0p8=0.7338 corrupt_frac_t_0p6_0p8=0.5574 out_w_norm=28.6215 out_g_norm=0.5405 acc_corrupt_t_0p0_0p2=0.1863 corrupt_frac_t_0p0_0p2=0.5474 acc_corrupt_t_0p8_1p0=0.9104 corrupt_frac_t_0p8_1p0=0.5585 acc_corrupt_t_0p2_0p4=0.3670 corrupt_frac_t_0p2_0p4=0.5583 acc_corrupt_t_0p4_0p6=0.5536 corrupt_frac_t_0p4_0p6=0.5594 loss_all=0.7114 init_gold_top10=0.7198 init_gold_top100=0.7198

LTA_openwebtext_dualt/logs/lta_owt_c1024_gpt2_cached_chunks_len1024_fast10k_4gpu_500step.log

@@ -0,0 +1,168 @@
+
+*****************************************
+Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+*****************************************
+[rank0]:[W512 16:35:00.260053068 ProcessGroupNCCL.cpp:4571] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.
+NCCL version 2.25.1+cuda12.8
+[rank1]:[W512 16:35:00.299852507 ProcessGroupNCCL.cpp:4571] [PG ID 0 PG GUID 0 Rank 1]  using GPU 1 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.
+[rank2]:[W512 16:35:00.301136581 ProcessGroupNCCL.cpp:4571] [PG ID 0 PG GUID 0 Rank 2]  using GPU 2 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.
+[rank3]:[W512 16:35:00.304399895 ProcessGroupNCCL.cpp:4571] [PG ID 0 PG GUID 0 Rank 3]  using GPU 3 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.
+{
+  "device": "cuda:0",
+  "rank": 0,
+  "world_size": 4,
+  "samples": "owt_cached_chunks:10904",
+  "vocab_size": 50257,
+  "save_dir": "runs/lta_owt_c1024_gpt2_cached_chunks_len1024_fast10k_4gpu_500step",
+  "batch_size": 32,
+  "grad_accum": 4,
+  "effective_batch_size": 512,
+  "global_batch_size": 512,
+  "lr_schedule": "constant_warmup",
+  "warmup_steps": 50,
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.999,
+  "adam_eps": 1e-08,
+  "model_type": "ddit",
+  "dual_t": true,
+  "corrupt_t_mode": "same",
+  "corrupt_min_t": 0.0,
+  "corrupt_max_t": 1.0,
+  "dirichlet_endpoint_mode": "categorical_dual_t",
+  "dirichlet_semantic_t_mode": "same",
+  "dirichlet_semantic_t_value": 0.0,
+  "categorical_wrong_from_full_vocab": true,
+  "simplex_bridge_sampler": "dirichlet",
+  "logistic_normal_sigma_min": 0.18,
+  "logistic_normal_sigma_max": 2.2,
+  "logistic_normal_tau_min": 0.65,
+  "logistic_normal_tau_max": 1.15,
+  "torch_compile": false,
+  "compile_mode": "max-autotune",
+  "state_format": "prob",
+  "target_loss": "hard_ce",
+  "meanflow_weight": 0.0,
+  "bridge_noise_init": "logistic_normal",
+  "noise_sigma": -1.0,
+  "wrap": true,
+  "wrap_mode": "stream",
+  "wrap_record_buffer_size": 200,
+  "owt_cached_chunks": true,
+  "owt_chunk_cache_dir": "/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext_lta_cached_chunks/gpt2_len1024_train_minus_100k_fast10k",
+  "owt_chunk_cache_rebuild": false,
+  "owt_chunk_cache_write_batch": 4096,
+  "online_chunk_shuffle": false,
+  "online_chunk_shuffle_buffer": 10000,
+  "openwebtext_split": "train_minus_100k",
+  "detokenizer": "auto",
+  "resolved_detokenizer": null,
+  "num_workers": 2,
+  "latest_every": 50,
+  "resume_path": ""
+}
+step=10 micro_steps=40 elapsed=67.7s lr=6.600000e-05 loss_all=10.7775 acc_all=0.6098 loss_corrupt=10.7889 acc_corrupt=0.4141 corrupt_frac=0.5619 loss=10.7889 loss_recon=10.7889 loss_meanflow=0.0000 mean_model_t=0.4986 mean_corrupt_t=0.4986 mean_loss_t_weight=1.0000 prior_center_loss_beta=0.0000 wrong_frac=0.4990 init_acc_corrupt=0.4668 init_gold_top10=0.4955 init_gold_top100=0.5245
+step=20 micro_steps=80 elapsed=57.9s lr=1.260000e-04 loss_all=10.3772 acc_all=0.3528 loss_corrupt=10.4082 acc_corrupt=0.2188 corrupt_frac=0.5554 loss=10.4082 loss_recon=10.4082 loss_meanflow=0.0000 mean_model_t=0.4975 mean_corrupt_t=0.4975 mean_loss_t_weight=1.0000 prior_center_loss_beta=0.0000 wrong_frac=0.4996 init_acc_corrupt=0.4680 init_gold_top10=0.4957 init_gold_top100=0.5226
+step=30 micro_steps=120 elapsed=59.3s lr=1.860000e-04 loss_all=9.3684 acc_all=0.2037 loss_corrupt=9.4072 acc_corrupt=0.1250 corrupt_frac=0.5514 loss=9.4072 loss_recon=9.4072 loss_meanflow=0.0000 mean_model_t=0.4946 mean_corrupt_t=0.4946 mean_loss_t_weight=1.0000 prior_center_loss_beta=0.0000 wrong_frac=0.5065 init_acc_corrupt=0.4578 init_gold_top10=0.4874 init_gold_top100=0.5184
+step=40 micro_steps=160 elapsed=62.9s lr=2.460000e-04 loss_all=8.1540 acc_all=0.2264 loss_corrupt=8.2181 acc_corrupt=0.1478 corrupt_frac=0.5382 loss=8.2181 loss_recon=8.2181 loss_meanflow=0.0000 mean_model_t=0.4880 mean_corrupt_t=0.4880 mean_loss_t_weight=1.0000 prior_center_loss_beta=0.0000 wrong_frac=0.5150 init_acc_corrupt=0.4486 init_gold_top10=0.4790 init_gold_top100=0.5111
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 797, in <module>
+[rank0]:     main()
+[rank0]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 693, in main
+[rank0]:     bridge = make_bridge(
+[rank0]:              ^^^^^^^^^^^^
+[rank0]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 313, in make_bridge
+[rank0]:     return make_dirichlet_bridge_batch(
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/flowtext_lab/bridges.py", line 529, in make_dirichlet_bridge_batch
+[rank0]:     state_probs = sample_dirichlet_bridge_simplex(
+[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/flowtext_lab/bridges.py", line 263, in sample_dirichlet_bridge_simplex
+[rank0]:     sample = sample.clamp_min(eps)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^
+[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.14 GiB. GPU 0 has a total capacity of 95.22 GiB of which 573.56 MiB is free. Process 971246 has 94.66 GiB memory in use. Of the allocated memory 61.51 GiB is allocated by PyTorch, and 31.86 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
+[rank1]: Traceback (most recent call last):
+[rank1]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 797, in <module>
+[rank1]:     main()
+[rank1]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 693, in main
+[rank1]:     bridge = make_bridge(
+[rank1]:              ^^^^^^^^^^^^
+[rank1]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 313, in make_bridge
+[rank1]:     return make_dirichlet_bridge_batch(
+[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank1]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/flowtext_lab/bridges.py", line 529, in make_dirichlet_bridge_batch
+[rank1]:     state_probs = sample_dirichlet_bridge_simplex(
+[rank1]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank1]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/flowtext_lab/bridges.py", line 263, in sample_dirichlet_bridge_simplex
+[rank1]:     sample = sample.clamp_min(eps)
+[rank1]:              ^^^^^^^^^^^^^^^^^^^^^
+[rank1]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.14 GiB. GPU 1 has a total capacity of 95.22 GiB of which 573.56 MiB is free. Process 971247 has 94.66 GiB memory in use. Of the allocated memory 61.51 GiB is allocated by PyTorch, and 31.86 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
+[rank3]: Traceback (most recent call last):
+[rank3]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 797, in <module>
+[rank3]:     main()
+[rank3]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 693, in main
+[rank3]:     bridge = make_bridge(
+[rank3]:              ^^^^^^^^^^^^
+[rank3]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 313, in make_bridge
+[rank3]:     return make_dirichlet_bridge_batch(
+[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank3]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/flowtext_lab/bridges.py", line 529, in make_dirichlet_bridge_batch
+[rank3]:     state_probs = sample_dirichlet_bridge_simplex(
+[rank3]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank3]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/flowtext_lab/bridges.py", line 263, in sample_dirichlet_bridge_simplex
+[rank3]:     sample = sample.clamp_min(eps)
+[rank3]:              ^^^^^^^^^^^^^^^^^^^^^
+[rank3]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.14 GiB. GPU 3 has a total capacity of 95.22 GiB of which 573.56 MiB is free. Process 971249 has 94.66 GiB memory in use. Of the allocated memory 61.51 GiB is allocated by PyTorch, and 31.86 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
+[rank2]: Traceback (most recent call last):
+[rank2]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 797, in <module>
+[rank2]:     main()
+[rank2]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 693, in main
+[rank2]:     bridge = make_bridge(
+[rank2]:              ^^^^^^^^^^^^
+[rank2]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/train.py", line 313, in make_bridge
+[rank2]:     return make_dirichlet_bridge_batch(
+[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank2]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/flowtext_lab/bridges.py", line 529, in make_dirichlet_bridge_batch
+[rank2]:     state_probs = sample_dirichlet_bridge_simplex(
+[rank2]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank2]:   File "/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/flowtext_lab/bridges.py", line 263, in sample_dirichlet_bridge_simplex
+[rank2]:     sample = sample.clamp_min(eps)
+[rank2]:              ^^^^^^^^^^^^^^^^^^^^^
+[rank2]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.14 GiB. GPU 2 has a total capacity of 95.22 GiB of which 573.56 MiB is free. Process 971248 has 94.66 GiB memory in use. Of the allocated memory 61.51 GiB is allocated by PyTorch, and 31.86 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
+[rank0]:[W512 16:40:20.480386792 ProcessGroupNCCL.cpp:1487] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
+W0512 16:40:21.089000 208015 torch/distributed/elastic/multiprocessing/api.py:898] Sending process 208083 closing signal SIGTERM
+W0512 16:40:21.090000 208015 torch/distributed/elastic/multiprocessing/api.py:898] Sending process 208084 closing signal SIGTERM
+W0512 16:40:21.090000 208015 torch/distributed/elastic/multiprocessing/api.py:898] Sending process 208085 closing signal SIGTERM
+E0512 16:40:21.469000 208015 torch/distributed/elastic/multiprocessing/api.py:870] failed (exitcode: 1) local_rank: 0 (pid: 208082) of binary: /usr/bin/python
+Traceback (most recent call last):
+  File "<frozen runpy>", line 198, in _run_module_as_main
+  File "<frozen runpy>", line 88, in _run_code
+  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/run.py", line 922, in <module>
+    main()
+  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
+    return f(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/run.py", line 918, in main
+    run(args)
+  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/run.py", line 909, in run
+    elastic_launch(
+  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/launcher/api.py", line 139, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/launcher/api.py", line 270, in launch_agent
+    raise ChildFailedError(
+torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
+============================================================
+train.py FAILED
+------------------------------------------------------------
+Failures:
+  <NO_OTHER_FAILURES>
+------------------------------------------------------------
+Root Cause (first observed failure):
+[0]:
+  time      : 2026-05-12_16:40:21
+  host      : localhost
+  rank      : 0 (local_rank: 0)
+  exitcode  : 1 (pid: 208082)
+  error_file: <N/A>
+  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
+============================================================

LTA_openwebtext_dualt/logs/lta_owt_c1024_len1024_lowt0p2_lowk80_noall_buf1000_gbs128_4gpu_100step.nohup.log

@@ -0,0 +1,73 @@
+[launch] owt low-t-cut low-k 100-step pilot
+[launch] run_name=lta_owt_c1024_len1024_lowt0p2_lowk80_noall_buf1000_gbs128_4gpu_100step
+[launch] save_dir=runs/lta_owt_c1024_len1024_lowt0p2_lowk80_noall_buf1000_gbs128_4gpu_100step
+[launch] t=0.20..1.0 mask=0.01..1.0
+[launch] mixture original=0.20 lowk=0.80 all=0.0
+
+*****************************************
+Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+*****************************************
+NCCL version 2.25.1+cuda12.8
+{
+  "device": "cuda:0",
+  "rank": 0,
+  "world_size": 4,
+  "samples": "wrapped_stream_online_shuffle:1000",
+  "vocab_size": 50257,
+  "save_dir": "runs/lta_owt_c1024_len1024_lowt0p2_lowk80_noall_buf1000_gbs128_4gpu_100step",
+  "batch_size": 16,
+  "grad_accum": 2,
+  "effective_batch_size": 128,
+  "global_batch_size": 128,
+  "lr_schedule": "constant_warmup",
+  "warmup_steps": 10,
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.999,
+  "adam_eps": 1e-08,
+  "model_type": "ddit",
+  "dual_t": true,
+  "corrupt_t_mode": "same",
+  "corrupt_min_t": 0.2,
+  "corrupt_max_t": 1.0,
+  "dirichlet_endpoint_mode": "categorical_dual_t",
+  "dirichlet_semantic_t_mode": "same",
+  "dirichlet_semantic_t_value": 0.0,
+  "categorical_wrong_from_full_vocab": true,
+  "simplex_bridge_sampler": "dirichlet",
+  "logistic_normal_sigma_min": 0.18,
+  "logistic_normal_sigma_max": 2.2,
+  "logistic_normal_tau_min": 0.65,
+  "logistic_normal_tau_max": 1.15,
+  "torch_compile": false,
+  "compile_mode": "max-autotune",
+  "state_format": "prob",
+  "target_loss": "hard_ce",
+  "meanflow_weight": 0.0,
+  "bridge_noise_init": "logistic_normal",
+  "noise_sigma": -1.0,
+  "wrap": true,
+  "wrap_mode": "stream",
+  "wrap_record_buffer_size": 200,
+  "owt_cached_chunks": false,
+  "owt_chunk_cache_dir": "",
+  "owt_chunk_cache_rebuild": false,
+  "owt_chunk_cache_write_batch": 4096,
+  "online_chunk_shuffle": true,
+  "online_chunk_shuffle_buffer": 1000,
+  "openwebtext_split": "train_minus_100k",
+  "detokenizer": "auto",
+  "resolved_detokenizer": null,
+  "num_workers": 0,
+  "latest_every": 50,
+  "resume_path": ""
+}
+step=10 micro_steps=20 elapsed=19.3s lr=3.000000e-04 loss_all=10.5503 acc_all=0.3705 loss_corrupt=10.5549 acc_corrupt=0.3380 corrupt_frac=0.8823 loss=10.5549 loss_recon=10.5549 loss_meanflow=0.0000 mean_model_t=0.5957 mean_corrupt_t=0.5957 wrong_frac=0.4038 init_acc_corrupt=0.5758 init_gold_top10=0.5953 init_gold_top100=0.6043
+step=20 micro_steps=40 elapsed=15.6s lr=3.000000e-04 loss_all=8.8827 acc_all=0.2839 loss_corrupt=8.9025 acc_corrupt=0.2544 corrupt_frac=0.8798 loss=8.9025 loss_recon=8.9025 loss_meanflow=0.0000 mean_model_t=0.6119 mean_corrupt_t=0.6119 wrong_frac=0.3848 init_acc_corrupt=0.5963 init_gold_top10=0.6144 init_gold_top100=0.6220
+step=30 micro_steps=60 elapsed=15.7s lr=3.000000e-04 loss_all=7.2594 acc_all=0.3284 loss_corrupt=7.3171 acc_corrupt=0.3054 corrupt_frac=0.8904 loss=7.3171 loss_recon=7.3171 loss_meanflow=0.0000 mean_model_t=0.6081 mean_corrupt_t=0.6081 wrong_frac=0.3883 init_acc_corrupt=0.5898 init_gold_top10=0.6111 init_gold_top100=0.6188
+step=40 micro_steps=80 elapsed=15.6s lr=3.000000e-04 loss_all=6.2331 acc_all=0.3068 loss_corrupt=6.4002 acc_corrupt=0.2793 corrupt_frac=0.8658 loss=6.4002 loss_recon=6.4002 loss_meanflow=0.0000 mean_model_t=0.5845 mean_corrupt_t=0.5845 wrong_frac=0.4127 init_acc_corrupt=0.5643 init_gold_top10=0.5865 init_gold_top100=0.5953
+step=50 micro_steps=100 elapsed=15.6s lr=3.000000e-04 loss_all=5.3938 acc_all=0.3656 loss_corrupt=5.6086 acc_corrupt=0.3343 corrupt_frac=0.8822 loss=5.6086 loss_recon=5.6086 loss_meanflow=0.0000 mean_model_t=0.6005 mean_corrupt_t=0.6005 wrong_frac=0.4020 init_acc_corrupt=0.5757 init_gold_top10=0.5974 init_gold_top100=0.6055
+step=60 micro_steps=120 elapsed=17.4s lr=3.000000e-04 loss_all=4.2746 acc_all=0.5067 loss_corrupt=4.5425 acc_corrupt=0.4699 corrupt_frac=0.8836 loss=4.5425 loss_recon=4.5425 loss_meanflow=0.0000 mean_model_t=0.6039 mean_corrupt_t=0.6039 wrong_frac=0.3956 init_acc_corrupt=0.5818 init_gold_top10=0.6036 init_gold_top100=0.6116
+step=70 micro_steps=140 elapsed=15.6s lr=3.000000e-04 loss_all=3.9509 acc_all=0.5396 loss_corrupt=4.2780 acc_corrupt=0.4999 corrupt_frac=0.8873 loss=4.2780 loss_recon=4.2780 loss_meanflow=0.0000 mean_model_t=0.5874 mean_corrupt_t=0.5874 wrong_frac=0.4177 init_acc_corrupt=0.5603 init_gold_top10=0.5815 init_gold_top100=0.5896
+step=80 micro_steps=160 elapsed=15.7s lr=3.000000e-04 loss_all=3.5467 acc_all=0.5825 loss_corrupt=3.9334 acc_corrupt=0.5363 corrupt_frac=0.8776 loss=3.9334 loss_recon=3.9334 loss_meanflow=0.0000 mean_model_t=0.5958 mean_corrupt_t=0.5958 wrong_frac=0.4048 init_acc_corrupt=0.5710 init_gold_top10=0.5943 init_gold_top100=0.6029
+step=90 micro_steps=180 elapsed=15.7s lr=3.000000e-04 loss_all=3.3378 acc_all=0.5975 loss_corrupt=3.7512 acc_corrupt=0.5472 corrupt_frac=0.8661 loss=3.7512 loss_recon=3.7512 loss_meanflow=0.0000 mean_model_t=0.5929 mean_corrupt_t=0.5929 wrong_frac=0.4062 init_acc_corrupt=0.5668 init_gold_top10=0.5927 init_gold_top100=0.6037
+step=100 micro_steps=200 elapsed=15.6s lr=3.000000e-04 loss_all=3.0738 acc_all=0.6320 loss_corrupt=3.4991 acc_corrupt=0.5806 corrupt_frac=0.8580 loss=3.4991 loss_recon=3.4991 loss_meanflow=0.0000 mean_model_t=0.6140 mean_corrupt_t=0.6140 wrong_frac=0.3817 init_acc_corrupt=0.5964 init_gold_top10=0.6177 init_gold_top100=0.6244

LTA_openwebtext_dualt/logs/lta_owt_launcher_opt_smoke_bspgpu64_gbs512_4gpu_20step_20260513.outer.log

@@ -0,0 +1,103 @@
+[launch] method=owt_categorical_fullvocab_c1024_fullycoupled host=di-20260411014000-djqhq time=2026-05-13T03:03:02+00:00
+[launch] cwd=/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+[launch] run_name=lta_owt_launcher_opt_smoke_bspgpu64_gbs512_4gpu_20step_20260513
+[launch] save_dir=runs/lta_owt_launcher_opt_smoke_bspgpu64_gbs512_4gpu_20step_20260513
+[launch] log_file=logs/lta_owt_launcher_opt_smoke_bspgpu64_gbs512_4gpu_20step_20260513.log
+[launch] data_path=/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext
+[launch] tokenizer=/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-standard/tokenizer.json
+[launch] split=train_minus_100k text_column=text
+[launch] nproc_per_node=4 global_batch_size=512 per_gpu_batch_size=64
+[launch] model d_model=768 n_layers=12 n_heads=12 dim_ff=3072 dropout=0.0
+[launch] optimizer=adamw lr=6e-4 wd=0.1 ema=0.0
+[launch] perf allow_tf32=1 activation_checkpointing=1 prefetch=4
+
+*****************************************
+Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+*****************************************
+NCCL version 2.25.1+cuda12.8
+{
+  "device": "cuda:0",
+  "rank": 0,
+  "world_size": 4,
+  "samples": "wrapped_stream",
+  "vocab_size": 50257,
+  "tokenizer_vocab_size": 50257,
+  "save_dir": "runs/lta_owt_launcher_opt_smoke_bspgpu64_gbs512_4gpu_20step_20260513",
+  "batch_size": 64,
+  "grad_accum": 2,
+  "effective_batch_size": 512,
+  "global_batch_size": 512,
+  "lr_schedule": "cosine",
+  "optimizer": "adamw",
+  "warmup_steps": 20,
+  "min_lr": 6e-05,
+  "weight_decay": 0.1,
+  "adamw_param_groups": "nanogpt",
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.95,
+  "adam_eps": 1e-08,
+  "muon_momentum": 0.95,
+  "muon_ns_steps": 5,
+  "muon_update_scale": 1.0,
+  "ema_decay": 0.0,
+  "ema_start_step": 0,
+  "model_type": "ddit",
+  "dual_t": true,
+  "corrupt_t_mode": "same",
+  "corrupt_min_t": 0.0,
+  "corrupt_max_t": 1.0,
+  "prefix_block_prob": 0.0,
+  "prefix_block_len": 128,
+  "dirichlet_endpoint_mode": "categorical_dual_t",
+  "dirichlet_semantic_t_mode": "same",
+  "dirichlet_semantic_t_value": 0.0,
+  "categorical_wrong_from_full_vocab": true,
+  "categorical_wrong_from_batch_valid_tokens": false,
+  "mask_mixture_original_prob": 0.0,
+  "mask_mixture_lowk_prob": 0.0,
+  "mask_mixture_lowcorrupt_prob": 0.0,
+  "mask_mixture_block_prob": 0.0,
+  "mask_mixture_all_prob": 0.0,
+  "mask_mixture_lowk_clean_tokens": "1,2,4,8,16,32,64",
+  "mask_mixture_lowcorrupt_tokens": "1,2,4,8,16,32,64",
+  "mask_mixture_block_tokens": "64,128",
+  "simplex_bridge_sampler": "dirichlet",
+  "logistic_normal_sigma_min": 0.18,
+  "logistic_normal_sigma_max": 2.2,
+  "logistic_normal_tau_min": 0.65,
+  "logistic_normal_tau_max": 1.15,
+  "torch_compile": false,
+  "compile_mode": "max-autotune",
+  "state_format": "prob",
+  "target_loss": "hard_ce",
+  "meanflow_weight": 0.0,
+  "bridge_noise_init": "logistic_normal",
+  "noise_sigma": -1.0,
+  "allow_tf32": true,
+  "activation_checkpointing": true,
+  "ddp_static_graph": false,
+  "ddp_gradient_as_bucket_view": true,
+  "blocking_data_transfer": false,
+  "dataloader_prefetch_factor": 4,
+  "full_train_stats": false,
+  "wrap": true,
+  "wrap_mode": "stream",
+  "wrap_record_buffer_size": 200,
+  "owt_cached_chunks": false,
+  "owt_chunk_cache_dir": "",
+  "owt_chunk_cache_rebuild": false,
+  "owt_chunk_cache_write_batch": 4096,
+  "owt_exact_repeat_per_chunk": 0,
+  "online_chunk_shuffle": false,
+  "online_chunk_shuffle_buffer": 10000,
+  "openwebtext_split": "train_minus_100k",
+  "detokenizer": "auto",
+  "resolved_detokenizer": null,
+  "num_workers": 4,
+  "latest_every": 0,
+  "resume_path": ""
+}
+step=5 micro_steps=10 elapsed=16.8s lr=1.800000e-04 loss_all=10.6893 acc_all=0.6116 loss_corrupt=10.7312 acc_corrupt=0.3232 corrupt_frac=0.5314 loss=10.7312 loss_recon=10.7312 loss_meanflow=0.0000 mean_model_t=0.5283 mean_corrupt_t=0.5283 mean_loss_t_weight=1.0000 prior_center_loss_beta=0.0000 wrong_frac=0.4484 init_acc_corrupt=0.3224 init_gold_top10=0.3393 init_gold_top100=0.3990
+step=10 micro_steps=20 elapsed=15.5s lr=3.300000e-04 loss_all=9.9390 acc_all=0.0729 loss_corrupt=9.9570 acc_corrupt=0.0432 corrupt_frac=0.5474 loss=9.9570 loss_recon=9.9570 loss_meanflow=0.0000 mean_model_t=0.5045 mean_corrupt_t=0.5045 mean_loss_t_weight=1.0000 prior_center_loss_beta=0.0000 wrong_frac=0.4738 init_acc_corrupt=0.3317 init_gold_top10=0.3498 init_gold_top100=0.3963
+step=15 micro_steps=30 elapsed=17.9s lr=4.800000e-04 loss_all=8.7640 acc_all=0.0341 loss_corrupt=8.7541 acc_corrupt=0.0343 corrupt_frac=0.6156 loss=8.7541 loss_recon=8.7541 loss_meanflow=0.0000 mean_model_t=0.5740 mean_corrupt_t=0.5740 mean_loss_t_weight=1.0000 prior_center_loss_beta=0.0000 wrong_frac=0.4372 init_acc_corrupt=0.3838 init_gold_top10=0.4053 init_gold_top100=0.4513
+Terminated

LTA_openwebtext_dualt/logs/owt_classic_fullvocab_len1024_infer_watch/infer_lta_owt_classic_fullvocab_bert_c1024_len1024_gbs512_8gpu_1m_save10k_20260521_162117_step_0010000_t1p45.log

@@ -0,0 +1,260 @@
+[watch-owt-classic] 2026-05-21_18:30:25 infer runs/lta_owt_classic_fullvocab_bert_c1024_len1024_gbs512_8gpu_1m_save10k_20260521_162117/step_0010000.pt -> docs/lta_samples/metrics_20260521/owt_classic_fullvocab_len1024_every10k_normal_steps_state_t1p45_c1024_n1024/lta_owt_classic_fullvocab_bert_c1024_len1024_gbs512_8gpu_1m_save10k_20260521_162117/step_0010000
+[ckpt] runs/lta_owt_classic_fullvocab_bert_c1024_len1024_gbs512_8gpu_1m_save10k_20260521_162117/step_0010000.pt step=10000
+[decode] steps128_c1024_t1p45 generated 4/1024
+[decode] steps128_c1024_t1p45 generated 8/1024
+[decode] steps128_c1024_t1p45 generated 12/1024
+[decode] steps128_c1024_t1p45 generated 16/1024
+[decode] steps128_c1024_t1p45 generated 20/1024
+[decode] steps128_c1024_t1p45 generated 24/1024
+[decode] steps128_c1024_t1p45 generated 28/1024
+[decode] steps128_c1024_t1p45 generated 32/1024
+[decode] steps128_c1024_t1p45 generated 36/1024
+[decode] steps128_c1024_t1p45 generated 40/1024
+[decode] steps128_c1024_t1p45 generated 44/1024
+[decode] steps128_c1024_t1p45 generated 48/1024
+[decode] steps128_c1024_t1p45 generated 52/1024
+[decode] steps128_c1024_t1p45 generated 56/1024
+[decode] steps128_c1024_t1p45 generated 60/1024
+[decode] steps128_c1024_t1p45 generated 64/1024
+[decode] steps128_c1024_t1p45 generated 68/1024
+[decode] steps128_c1024_t1p45 generated 72/1024
+[decode] steps128_c1024_t1p45 generated 76/1024
+[decode] steps128_c1024_t1p45 generated 80/1024
+[decode] steps128_c1024_t1p45 generated 84/1024
+[decode] steps128_c1024_t1p45 generated 88/1024
+[decode] steps128_c1024_t1p45 generated 92/1024
+[decode] steps128_c1024_t1p45 generated 96/1024
+[decode] steps128_c1024_t1p45 generated 100/1024
+[decode] steps128_c1024_t1p45 generated 104/1024
+[decode] steps128_c1024_t1p45 generated 108/1024
+[decode] steps128_c1024_t1p45 generated 112/1024
+[decode] steps128_c1024_t1p45 generated 116/1024
+[decode] steps128_c1024_t1p45 generated 120/1024
+[decode] steps128_c1024_t1p45 generated 124/1024
+[decode] steps128_c1024_t1p45 generated 128/1024
+[decode] steps128_c1024_t1p45 generated 132/1024
+[decode] steps128_c1024_t1p45 generated 136/1024
+[decode] steps128_c1024_t1p45 generated 140/1024
+[decode] steps128_c1024_t1p45 generated 144/1024
+[decode] steps128_c1024_t1p45 generated 148/1024
+[decode] steps128_c1024_t1p45 generated 152/1024
+[decode] steps128_c1024_t1p45 generated 156/1024
+[decode] steps128_c1024_t1p45 generated 160/1024
+[decode] steps128_c1024_t1p45 generated 164/1024
+[decode] steps128_c1024_t1p45 generated 168/1024
+[decode] steps128_c1024_t1p45 generated 172/1024
+[decode] steps128_c1024_t1p45 generated 176/1024
+[decode] steps128_c1024_t1p45 generated 180/1024
+[decode] steps128_c1024_t1p45 generated 184/1024
+[decode] steps128_c1024_t1p45 generated 188/1024
+[decode] steps128_c1024_t1p45 generated 192/1024
+[decode] steps128_c1024_t1p45 generated 196/1024
+[decode] steps128_c1024_t1p45 generated 200/1024
+[decode] steps128_c1024_t1p45 generated 204/1024
+[decode] steps128_c1024_t1p45 generated 208/1024
+[decode] steps128_c1024_t1p45 generated 212/1024
+[decode] steps128_c1024_t1p45 generated 216/1024
+[decode] steps128_c1024_t1p45 generated 220/1024
+[decode] steps128_c1024_t1p45 generated 224/1024
+[decode] steps128_c1024_t1p45 generated 228/1024
+[decode] steps128_c1024_t1p45 generated 232/1024
+[decode] steps128_c1024_t1p45 generated 236/1024
+[decode] steps128_c1024_t1p45 generated 240/1024
+[decode] steps128_c1024_t1p45 generated 244/1024
+[decode] steps128_c1024_t1p45 generated 248/1024
+[decode] steps128_c1024_t1p45 generated 252/1024
+[decode] steps128_c1024_t1p45 generated 256/1024
+[decode] steps128_c1024_t1p45 generated 260/1024
+[decode] steps128_c1024_t1p45 generated 264/1024
+[decode] steps128_c1024_t1p45 generated 268/1024
+[decode] steps128_c1024_t1p45 generated 272/1024
+[decode] steps128_c1024_t1p45 generated 276/1024
+[decode] steps128_c1024_t1p45 generated 280/1024
+[decode] steps128_c1024_t1p45 generated 284/1024
+[decode] steps128_c1024_t1p45 generated 288/1024
+[decode] steps128_c1024_t1p45 generated 292/1024
+[decode] steps128_c1024_t1p45 generated 296/1024
+[decode] steps128_c1024_t1p45 generated 300/1024
+[decode] steps128_c1024_t1p45 generated 304/1024
+[decode] steps128_c1024_t1p45 generated 308/1024
+[decode] steps128_c1024_t1p45 generated 312/1024
+[decode] steps128_c1024_t1p45 generated 316/1024
+[decode] steps128_c1024_t1p45 generated 320/1024
+[decode] steps128_c1024_t1p45 generated 324/1024
+[decode] steps128_c1024_t1p45 generated 328/1024
+[decode] steps128_c1024_t1p45 generated 332/1024
+[decode] steps128_c1024_t1p45 generated 336/1024
+[decode] steps128_c1024_t1p45 generated 340/1024
+[decode] steps128_c1024_t1p45 generated 344/1024
+[decode] steps128_c1024_t1p45 generated 348/1024
+[decode] steps128_c1024_t1p45 generated 352/1024
+[decode] steps128_c1024_t1p45 generated 356/1024
+[decode] steps128_c1024_t1p45 generated 360/1024
+[decode] steps128_c1024_t1p45 generated 364/1024
+[decode] steps128_c1024_t1p45 generated 368/1024
+[decode] steps128_c1024_t1p45 generated 372/1024
+[decode] steps128_c1024_t1p45 generated 376/1024
+[decode] steps128_c1024_t1p45 generated 380/1024
+[decode] steps128_c1024_t1p45 generated 384/1024
+[decode] steps128_c1024_t1p45 generated 388/1024
+[decode] steps128_c1024_t1p45 generated 392/1024
+[decode] steps128_c1024_t1p45 generated 396/1024
+[decode] steps128_c1024_t1p45 generated 400/1024
+[decode] steps128_c1024_t1p45 generated 404/1024
+[decode] steps128_c1024_t1p45 generated 408/1024
+[decode] steps128_c1024_t1p45 generated 412/1024
+[decode] steps128_c1024_t1p45 generated 416/1024
+[decode] steps128_c1024_t1p45 generated 420/1024
+[decode] steps128_c1024_t1p45 generated 424/1024
+[decode] steps128_c1024_t1p45 generated 428/1024
+[decode] steps128_c1024_t1p45 generated 432/1024
+[decode] steps128_c1024_t1p45 generated 436/1024
+[decode] steps128_c1024_t1p45 generated 440/1024
+[decode] steps128_c1024_t1p45 generated 444/1024
+[decode] steps128_c1024_t1p45 generated 448/1024
+[decode] steps128_c1024_t1p45 generated 452/1024
+[decode] steps128_c1024_t1p45 generated 456/1024
+[decode] steps128_c1024_t1p45 generated 460/1024
+[decode] steps128_c1024_t1p45 generated 464/1024
+[decode] steps128_c1024_t1p45 generated 468/1024
+[decode] steps128_c1024_t1p45 generated 472/1024
+[decode] steps128_c1024_t1p45 generated 476/1024
+[decode] steps128_c1024_t1p45 generated 480/1024
+[decode] steps128_c1024_t1p45 generated 484/1024
+[decode] steps128_c1024_t1p45 generated 488/1024
+[decode] steps128_c1024_t1p45 generated 492/1024
+[decode] steps128_c1024_t1p45 generated 496/1024
+[decode] steps128_c1024_t1p45 generated 500/1024
+[decode] steps128_c1024_t1p45 generated 504/1024
+[decode] steps128_c1024_t1p45 generated 508/1024
+[decode] steps128_c1024_t1p45 generated 512/1024
+[decode] steps128_c1024_t1p45 generated 516/1024
+[decode] steps128_c1024_t1p45 generated 520/1024
+[decode] steps128_c1024_t1p45 generated 524/1024
+[decode] steps128_c1024_t1p45 generated 528/1024
+[decode] steps128_c1024_t1p45 generated 532/1024
+[decode] steps128_c1024_t1p45 generated 536/1024
+[decode] steps128_c1024_t1p45 generated 540/1024
+[decode] steps128_c1024_t1p45 generated 544/1024
+[decode] steps128_c1024_t1p45 generated 548/1024
+[decode] steps128_c1024_t1p45 generated 552/1024
+[decode] steps128_c1024_t1p45 generated 556/1024
+[decode] steps128_c1024_t1p45 generated 560/1024
+[decode] steps128_c1024_t1p45 generated 564/1024
+[decode] steps128_c1024_t1p45 generated 568/1024
+[decode] steps128_c1024_t1p45 generated 572/1024
+[decode] steps128_c1024_t1p45 generated 576/1024
+[decode] steps128_c1024_t1p45 generated 580/1024
+[decode] steps128_c1024_t1p45 generated 584/1024
+[decode] steps128_c1024_t1p45 generated 588/1024
+[decode] steps128_c1024_t1p45 generated 592/1024
+[decode] steps128_c1024_t1p45 generated 596/1024
+[decode] steps128_c1024_t1p45 generated 600/1024
+[decode] steps128_c1024_t1p45 generated 604/1024
+[decode] steps128_c1024_t1p45 generated 608/1024
+[decode] steps128_c1024_t1p45 generated 612/1024
+[decode] steps128_c1024_t1p45 generated 616/1024
+[decode] steps128_c1024_t1p45 generated 620/1024
+[decode] steps128_c1024_t1p45 generated 624/1024
+[decode] steps128_c1024_t1p45 generated 628/1024
+[decode] steps128_c1024_t1p45 generated 632/1024
+[decode] steps128_c1024_t1p45 generated 636/1024
+[decode] steps128_c1024_t1p45 generated 640/1024
+[decode] steps128_c1024_t1p45 generated 644/1024
+[decode] steps128_c1024_t1p45 generated 648/1024
+[decode] steps128_c1024_t1p45 generated 652/1024
+[decode] steps128_c1024_t1p45 generated 656/1024
+[decode] steps128_c1024_t1p45 generated 660/1024
+[decode] steps128_c1024_t1p45 generated 664/1024
+[decode] steps128_c1024_t1p45 generated 668/1024
+[decode] steps128_c1024_t1p45 generated 672/1024
+[decode] steps128_c1024_t1p45 generated 676/1024
+[decode] steps128_c1024_t1p45 generated 680/1024
+[decode] steps128_c1024_t1p45 generated 684/1024
+[decode] steps128_c1024_t1p45 generated 688/1024
+[decode] steps128_c1024_t1p45 generated 692/1024
+[decode] steps128_c1024_t1p45 generated 696/1024
+[decode] steps128_c1024_t1p45 generated 700/1024
+[decode] steps128_c1024_t1p45 generated 704/1024
+[decode] steps128_c1024_t1p45 generated 708/1024
+[decode] steps128_c1024_t1p45 generated 712/1024
+[decode] steps128_c1024_t1p45 generated 716/1024
+[decode] steps128_c1024_t1p45 generated 720/1024
+[decode] steps128_c1024_t1p45 generated 724/1024
+[decode] steps128_c1024_t1p45 generated 728/1024
+[decode] steps128_c1024_t1p45 generated 732/1024
+[decode] steps128_c1024_t1p45 generated 736/1024
+[decode] steps128_c1024_t1p45 generated 740/1024
+[decode] steps128_c1024_t1p45 generated 744/1024
+[decode] steps128_c1024_t1p45 generated 748/1024
+[decode] steps128_c1024_t1p45 generated 752/1024
+[decode] steps128_c1024_t1p45 generated 756/1024
+[decode] steps128_c1024_t1p45 generated 760/1024
+[decode] steps128_c1024_t1p45 generated 764/1024
+[decode] steps128_c1024_t1p45 generated 768/1024
+[decode] steps128_c1024_t1p45 generated 772/1024
+[decode] steps128_c1024_t1p45 generated 776/1024
+[decode] steps128_c1024_t1p45 generated 780/1024
+[decode] steps128_c1024_t1p45 generated 784/1024
+[decode] steps128_c1024_t1p45 generated 788/1024
+[decode] steps128_c1024_t1p45 generated 792/1024
+[decode] steps128_c1024_t1p45 generated 796/1024
+[decode] steps128_c1024_t1p45 generated 800/1024
+[decode] steps128_c1024_t1p45 generated 804/1024
+[decode] steps128_c1024_t1p45 generated 808/1024
+[decode] steps128_c1024_t1p45 generated 812/1024
+[decode] steps128_c1024_t1p45 generated 816/1024
+[decode] steps128_c1024_t1p45 generated 820/1024
+[decode] steps128_c1024_t1p45 generated 824/1024
+[decode] steps128_c1024_t1p45 generated 828/1024
+[decode] steps128_c1024_t1p45 generated 832/1024
+[decode] steps128_c1024_t1p45 generated 836/1024
+[decode] steps128_c1024_t1p45 generated 840/1024
+[decode] steps128_c1024_t1p45 generated 844/1024
+[decode] steps128_c1024_t1p45 generated 848/1024
+[decode] steps128_c1024_t1p45 generated 852/1024
+[decode] steps128_c1024_t1p45 generated 856/1024
+[decode] steps128_c1024_t1p45 generated 860/1024
+[decode] steps128_c1024_t1p45 generated 864/1024
+[decode] steps128_c1024_t1p45 generated 868/1024
+[decode] steps128_c1024_t1p45 generated 872/1024
+[decode] steps128_c1024_t1p45 generated 876/1024
+[decode] steps128_c1024_t1p45 generated 880/1024
+[decode] steps128_c1024_t1p45 generated 884/1024
+[decode] steps128_c1024_t1p45 generated 888/1024
+[decode] steps128_c1024_t1p45 generated 892/1024
+[decode] steps128_c1024_t1p45 generated 896/1024
+[decode] steps128_c1024_t1p45 generated 900/1024
+[decode] steps128_c1024_t1p45 generated 904/1024
+[decode] steps128_c1024_t1p45 generated 908/1024
+[decode] steps128_c1024_t1p45 generated 912/1024
+[decode] steps128_c1024_t1p45 generated 916/1024
+[decode] steps128_c1024_t1p45 generated 920/1024
+[decode] steps128_c1024_t1p45 generated 924/1024
+[decode] steps128_c1024_t1p45 generated 928/1024
+[decode] steps128_c1024_t1p45 generated 932/1024
+[decode] steps128_c1024_t1p45 generated 936/1024
+[decode] steps128_c1024_t1p45 generated 940/1024
+[decode] steps128_c1024_t1p45 generated 944/1024
+[decode] steps128_c1024_t1p45 generated 948/1024
+[decode] steps128_c1024_t1p45 generated 952/1024
+[decode] steps128_c1024_t1p45 generated 956/1024
+[decode] steps128_c1024_t1p45 generated 960/1024
+[decode] steps128_c1024_t1p45 generated 964/1024
+[decode] steps128_c1024_t1p45 generated 968/1024
+[decode] steps128_c1024_t1p45 generated 972/1024
+[decode] steps128_c1024_t1p45 generated 976/1024
+[decode] steps128_c1024_t1p45 generated 980/1024
+[decode] steps128_c1024_t1p45 generated 984/1024
+[decode] steps128_c1024_t1p45 generated 988/1024
+[decode] steps128_c1024_t1p45 generated 992/1024
+[decode] steps128_c1024_t1p45 generated 996/1024
+[decode] steps128_c1024_t1p45 generated 1000/1024
+[decode] steps128_c1024_t1p45 generated 1004/1024
+[decode] steps128_c1024_t1p45 generated 1008/1024
+[decode] steps128_c1024_t1p45 generated 1012/1024
+[decode] steps128_c1024_t1p45 generated 1016/1024
+[decode] steps128_c1024_t1p45 generated 1020/1024
+[decode] steps128_c1024_t1p45 generated 1024/1024
+[summary] {"name": "steps128_c1024_t1p45", "step": 10000, "decode_steps": 128, "concentration_max": 1024.0, "raw_genppl": 3.7899707881267246, "stripped_genppl": 3.7899707881267246, "sample_entropy": 0.9235058106966801, "distinct_1": 0.0008497238159179688, "distinct_2": 0.002539253421309873, "top_token_mass": 0.5727787017822266, "raw_kept": 1024, "stripped_kept": 1024}
+[watch-owt-classic] 2026-05-21_19:13:52 done step_0010000

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/random/__init__.pxd

@@ -0,0 +1,14 @@
+cimport numpy as np
+from libc.stdint cimport uint32_t, uint64_t
+
+cdef extern from "numpy/random/bitgen.h":
+    struct bitgen:
+        void *state
+        uint64_t (*next_uint64)(void *st) nogil
+        uint32_t (*next_uint32)(void *st) nogil
+        double (*next_double)(void *st) nogil
+        uint64_t (*next_raw)(void *st) nogil
+
+    ctypedef bitgen bitgen_t
+
+from numpy.random.bit_generator cimport BitGenerator, SeedSequence

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/random/__init__.py

@@ -0,0 +1,215 @@
+"""
+========================
+Random Number Generation
+========================
+
+Use ``default_rng()`` to create a `Generator` and call its methods.
+
+=============== =========================================================
+Generator
+--------------- ---------------------------------------------------------
+Generator       Class implementing all of the random number distributions
+default_rng     Default constructor for ``Generator``
+=============== =========================================================
+
+============================================= ===
+BitGenerator Streams that work with Generator
+--------------------------------------------- ---
+MT19937
+PCG64
+PCG64DXSM
+Philox
+SFC64
+============================================= ===
+
+============================================= ===
+Getting entropy to initialize a BitGenerator
+--------------------------------------------- ---
+SeedSequence
+============================================= ===
+
+
+Legacy
+------
+
+For backwards compatibility with previous versions of numpy before 1.17, the
+various aliases to the global `RandomState` methods are left alone and do not
+use the new `Generator` API.
+
+==================== =========================================================
+Utility functions
+-------------------- ---------------------------------------------------------
+random               Uniformly distributed floats over ``[0, 1)``
+bytes                Uniformly distributed random bytes.
+permutation          Randomly permute a sequence / generate a random sequence.
+shuffle              Randomly permute a sequence in place.
+choice               Random sample from 1-D array.
+==================== =========================================================
+
+==================== =========================================================
+Compatibility
+functions - removed
+in the new API
+-------------------- ---------------------------------------------------------
+rand                 Uniformly distributed values.
+randn                Normally distributed values.
+ranf                 Uniformly distributed floating point numbers.
+random_integers      Uniformly distributed integers in a given range.
+                     (deprecated, use ``integers(..., closed=True)`` instead)
+random_sample        Alias for `random_sample`
+randint              Uniformly distributed integers in a given range
+seed                 Seed the legacy random number generator.
+==================== =========================================================
+
+==================== =========================================================
+Univariate
+distributions
+-------------------- ---------------------------------------------------------
+beta                 Beta distribution over ``[0, 1]``.
+binomial             Binomial distribution.
+chisquare            :math:`\\chi^2` distribution.
+exponential          Exponential distribution.
+f                    F (Fisher-Snedecor) distribution.
+gamma                Gamma distribution.
+geometric            Geometric distribution.
+gumbel               Gumbel distribution.
+hypergeometric       Hypergeometric distribution.
+laplace              Laplace distribution.
+logistic             Logistic distribution.
+lognormal            Log-normal distribution.
+logseries            Logarithmic series distribution.
+negative_binomial    Negative binomial distribution.
+noncentral_chisquare Non-central chi-square distribution.
+noncentral_f         Non-central F distribution.
+normal               Normal / Gaussian distribution.
+pareto               Pareto distribution.
+poisson              Poisson distribution.
+power                Power distribution.
+rayleigh             Rayleigh distribution.
+triangular           Triangular distribution.
+uniform              Uniform distribution.
+vonmises             Von Mises circular distribution.
+wald                 Wald (inverse Gaussian) distribution.
+weibull              Weibull distribution.
+zipf                 Zipf's distribution over ranked data.
+==================== =========================================================
+
+==================== ==========================================================
+Multivariate
+distributions
+-------------------- ----------------------------------------------------------
+dirichlet            Multivariate generalization of Beta distribution.
+multinomial          Multivariate generalization of the binomial distribution.
+multivariate_normal  Multivariate generalization of the normal distribution.
+==================== ==========================================================
+
+==================== =========================================================
+Standard
+distributions
+-------------------- ---------------------------------------------------------
+standard_cauchy      Standard Cauchy-Lorentz distribution.
+standard_exponential Standard exponential distribution.
+standard_gamma       Standard Gamma distribution.
+standard_normal      Standard normal distribution.
+standard_t           Standard Student's t-distribution.
+==================== =========================================================
+
+==================== =========================================================
+Internal functions
+-------------------- ---------------------------------------------------------
+get_state            Get tuple representing internal state of generator.
+set_state            Set state of generator.
+==================== =========================================================
+
+
+"""
+__all__ = [
+    'beta',
+    'binomial',
+    'bytes',
+    'chisquare',
+    'choice',
+    'dirichlet',
+    'exponential',
+    'f',
+    'gamma',
+    'geometric',
+    'get_state',
+    'gumbel',
+    'hypergeometric',
+    'laplace',
+    'logistic',
+    'lognormal',
+    'logseries',
+    'multinomial',
+    'multivariate_normal',
+    'negative_binomial',
+    'noncentral_chisquare',
+    'noncentral_f',
+    'normal',
+    'pareto',
+    'permutation',
+    'poisson',
+    'power',
+    'rand',
+    'randint',
+    'randn',
+    'random',
+    'random_integers',
+    'random_sample',
+    'ranf',
+    'rayleigh',
+    'sample',
+    'seed',
+    'set_state',
+    'shuffle',
+    'standard_cauchy',
+    'standard_exponential',
+    'standard_gamma',
+    'standard_normal',
+    'standard_t',
+    'triangular',
+    'uniform',
+    'vonmises',
+    'wald',
+    'weibull',
+    'zipf',
+]
+
+# add these for module-freeze analysis (like PyInstaller)
+from . import _pickle
+from . import _common
+from . import _bounded_integers
+
+from ._generator import Generator, default_rng
+from .bit_generator import SeedSequence, BitGenerator
+from ._mt19937 import MT19937
+from ._pcg64 import PCG64, PCG64DXSM
+from ._philox import Philox
+from ._sfc64 import SFC64
+from .mtrand import *
+
+__all__ += ['Generator', 'RandomState', 'SeedSequence', 'MT19937',
+            'Philox', 'PCG64', 'PCG64DXSM', 'SFC64', 'default_rng',
+            'BitGenerator']
+
+
+def __RandomState_ctor():
+    """Return a RandomState instance.
+
+    This function exists solely to assist (un)pickling.
+
+    Note that the state of the RandomState returned here is irrelevant, as this
+    function's entire purpose is to return a newly allocated RandomState whose
+    state pickle can set.  Consequently the RandomState returned by this function
+    is a freshly allocated copy with a seed=0.
+
+    See https://github.com/numpy/numpy/issues/4763 for a detailed discussion
+
+    """
+    return RandomState(seed=0)
+
+
+from numpy._pytesttester import PytestTester
+test = PytestTester(__name__)
+del PytestTester

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/random/__init__.pyi

@@ -0,0 +1,72 @@
+from numpy._pytesttester import PytestTester
+
+from numpy.random._generator import Generator as Generator
+from numpy.random._generator import default_rng as default_rng
+from numpy.random._mt19937 import MT19937 as MT19937
+from numpy.random._pcg64 import (
+    PCG64 as PCG64,
+    PCG64DXSM as PCG64DXSM,
+)
+from numpy.random._philox import Philox as Philox
+from numpy.random._sfc64 import SFC64 as SFC64
+from numpy.random.bit_generator import BitGenerator as BitGenerator
+from numpy.random.bit_generator import SeedSequence as SeedSequence
+from numpy.random.mtrand import (
+    RandomState as RandomState,
+    beta as beta,
+    binomial as binomial,
+    bytes as bytes,
+    chisquare as chisquare,
+    choice as choice,
+    dirichlet as dirichlet,
+    exponential as exponential,
+    f as f,
+    gamma as gamma,
+    geometric as geometric,
+    get_bit_generator as get_bit_generator,
+    get_state as get_state,
+    gumbel as gumbel,
+    hypergeometric as hypergeometric,
+    laplace as laplace,
+    logistic as logistic,
+    lognormal as lognormal,
+    logseries as logseries,
+    multinomial as multinomial,
+    multivariate_normal as multivariate_normal,
+    negative_binomial as negative_binomial,
+    noncentral_chisquare as noncentral_chisquare,
+    noncentral_f as noncentral_f,
+    normal as normal,
+    pareto as pareto,
+    permutation as permutation,
+    poisson as poisson,
+    power as power,
+    rand as rand,
+    randint as randint,
+    randn as randn,
+    random as random,
+    random_integers as random_integers,
+    random_sample as random_sample,
+    ranf as ranf,
+    rayleigh as rayleigh,
+    sample as sample,
+    seed as seed,
+    set_bit_generator as set_bit_generator,
+    set_state as set_state,
+    shuffle as shuffle,
+    standard_cauchy as standard_cauchy,
+    standard_exponential as standard_exponential,
+    standard_gamma as standard_gamma,
+    standard_normal as standard_normal,
+    standard_t as standard_t,
+    triangular as triangular,
+    uniform as uniform,
+    vonmises as vonmises,
+    wald as wald,
+    weibull as weibull,
+    zipf as zipf,
+)
+
+__all__: list[str]
+__path__: list[str]
+test: PytestTester

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/random/_common.pxd

@@ -0,0 +1,106 @@
+#cython: language_level=3
+
+from libc.stdint cimport uint32_t, uint64_t, int32_t, int64_t
+
+import numpy as np
+cimport numpy as np
+
+from numpy.random cimport bitgen_t
+
+cdef double POISSON_LAM_MAX
+cdef double LEGACY_POISSON_LAM_MAX
+cdef uint64_t MAXSIZE
+
+cdef enum ConstraintType:
+    CONS_NONE
+    CONS_NON_NEGATIVE
+    CONS_POSITIVE
+    CONS_POSITIVE_NOT_NAN
+    CONS_BOUNDED_0_1
+    CONS_BOUNDED_GT_0_1
+    CONS_BOUNDED_LT_0_1
+    CONS_GT_1
+    CONS_GTE_1
+    CONS_POISSON
+    LEGACY_CONS_POISSON
+
+ctypedef ConstraintType constraint_type
+
+cdef object benchmark(bitgen_t *bitgen, object lock, Py_ssize_t cnt, object method)
+cdef object random_raw(bitgen_t *bitgen, object lock, object size, object output)
+cdef object prepare_cffi(bitgen_t *bitgen)
+cdef object prepare_ctypes(bitgen_t *bitgen)
+cdef int check_constraint(double val, object name, constraint_type cons) except -1
+cdef int check_array_constraint(np.ndarray val, object name, constraint_type cons) except -1
+
+cdef extern from "include/aligned_malloc.h":
+    cdef void *PyArray_realloc_aligned(void *p, size_t n)
+    cdef void *PyArray_malloc_aligned(size_t n)
+    cdef void *PyArray_calloc_aligned(size_t n, size_t s)
+    cdef void PyArray_free_aligned(void *p)
+
+ctypedef void (*random_double_fill)(bitgen_t *state, np.npy_intp count, double* out)  noexcept nogil
+ctypedef double (*random_double_0)(void *state)  noexcept nogil
+ctypedef double (*random_double_1)(void *state, double a)  noexcept nogil
+ctypedef double (*random_double_2)(void *state, double a, double b)  noexcept nogil
+ctypedef double (*random_double_3)(void *state, double a, double b, double c)  noexcept nogil
+
+ctypedef void (*random_float_fill)(bitgen_t *state, np.npy_intp count, float* out)  noexcept nogil
+ctypedef float (*random_float_0)(bitgen_t *state)  noexcept nogil
+ctypedef float (*random_float_1)(bitgen_t *state, float a)  noexcept nogil
+
+ctypedef int64_t (*random_uint_0)(void *state)  noexcept nogil
+ctypedef int64_t (*random_uint_d)(void *state, double a)  noexcept nogil
+ctypedef int64_t (*random_uint_dd)(void *state, double a, double b)  noexcept nogil
+ctypedef int64_t (*random_uint_di)(void *state, double a, uint64_t b)  noexcept nogil
+ctypedef int64_t (*random_uint_i)(void *state, int64_t a)  noexcept nogil
+ctypedef int64_t (*random_uint_iii)(void *state, int64_t a, int64_t b, int64_t c)  noexcept nogil
+
+ctypedef uint32_t (*random_uint_0_32)(bitgen_t *state)  noexcept nogil
+ctypedef uint32_t (*random_uint_1_i_32)(bitgen_t *state, uint32_t a)  noexcept nogil
+
+ctypedef int32_t (*random_int_2_i_32)(bitgen_t *state, int32_t a, int32_t b)  noexcept nogil
+ctypedef int64_t (*random_int_2_i)(bitgen_t *state, int64_t a, int64_t b)  noexcept nogil
+
+cdef double kahan_sum(double *darr, np.npy_intp n) noexcept
+
+cdef inline double uint64_to_double(uint64_t rnd) noexcept nogil:
+    return (rnd >> 11) * (1.0 / 9007199254740992.0)
+
+cdef object double_fill(void *func, bitgen_t *state, object size, object lock, object out)
+
+cdef object float_fill(void *func, bitgen_t *state, object size, object lock, object out)
+
+cdef object float_fill_from_double(void *func, bitgen_t *state, object size, object lock, object out)
+
+cdef object wrap_int(object val, object bits)
+
+cdef np.ndarray int_to_array(object value, object name, object bits, object uint_size)
+
+cdef validate_output_shape(iter_shape, np.ndarray output)
+
+cdef object cont(void *func, void *state, object size, object lock, int narg,
+                 object a, object a_name, constraint_type a_constraint,
+                 object b, object b_name, constraint_type b_constraint,
+                 object c, object c_name, constraint_type c_constraint,
+                 object out)
+
+cdef object disc(void *func, void *state, object size, object lock,
+                 int narg_double, int narg_int64,
+                 object a, object a_name, constraint_type a_constraint,
+                 object b, object b_name, constraint_type b_constraint,
+                 object c, object c_name, constraint_type c_constraint)
+
+cdef object cont_f(void *func, bitgen_t *state, object size, object lock,
+                   object a, object a_name, constraint_type a_constraint,
+                   object out)
+
+cdef object cont_broadcast_3(void *func, void *state, object size, object lock,
+                             np.ndarray a_arr, object a_name, constraint_type a_constraint,
+                             np.ndarray b_arr, object b_name, constraint_type b_constraint,
+                             np.ndarray c_arr, object c_name, constraint_type c_constraint)
+
+cdef object discrete_broadcast_iii(void *func, void *state, object size, object lock,
+                                   np.ndarray a_arr, object a_name, constraint_type a_constraint,
+                                   np.ndarray b_arr, object b_name, constraint_type b_constraint,
+                                   np.ndarray c_arr, object c_name, constraint_type c_constraint)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/random/_generator.pyi

@@ -0,0 +1,681 @@
+from collections.abc import Callable
+from typing import Any, Union, overload, TypeVar, Literal
+
+from numpy import (
+    bool_,
+    dtype,
+    float32,
+    float64,
+    int8,
+    int16,
+    int32,
+    int64,
+    int_,
+    ndarray,
+    uint,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+)
+from numpy.random import BitGenerator, SeedSequence
+from numpy._typing import (
+    ArrayLike,
+    _ArrayLikeFloat_co,
+    _ArrayLikeInt_co,
+    _DoubleCodes,
+    _DTypeLikeBool,
+    _DTypeLikeInt,
+    _DTypeLikeUInt,
+    _Float32Codes,
+    _Float64Codes,
+    _FloatLike_co,
+    _Int8Codes,
+    _Int16Codes,
+    _Int32Codes,
+    _Int64Codes,
+    _IntCodes,
+    _ShapeLike,
+    _SingleCodes,
+    _SupportsDType,
+    _UInt8Codes,
+    _UInt16Codes,
+    _UInt32Codes,
+    _UInt64Codes,
+    _UIntCodes,
+)
+
+_ArrayType = TypeVar("_ArrayType", bound=ndarray[Any, Any])
+
+_DTypeLikeFloat32 = Union[
+    dtype[float32],
+    _SupportsDType[dtype[float32]],
+    type[float32],
+    _Float32Codes,
+    _SingleCodes,
+]
+
+_DTypeLikeFloat64 = Union[
+    dtype[float64],
+    _SupportsDType[dtype[float64]],
+    type[float],
+    type[float64],
+    _Float64Codes,
+    _DoubleCodes,
+]
+
+class Generator:
+    def __init__(self, bit_generator: BitGenerator) -> None: ...
+    def __repr__(self) -> str: ...
+    def __str__(self) -> str: ...
+    def __getstate__(self) -> dict[str, Any]: ...
+    def __setstate__(self, state: dict[str, Any]) -> None: ...
+    def __reduce__(self) -> tuple[Callable[[str], Generator], tuple[str], dict[str, Any]]: ...
+    @property
+    def bit_generator(self) -> BitGenerator: ...
+    def spawn(self, n_children: int) -> list[Generator]: ...
+    def bytes(self, length: int) -> bytes: ...
+    @overload
+    def standard_normal(  # type: ignore[misc]
+        self,
+        size: None = ...,
+        dtype: _DTypeLikeFloat32 | _DTypeLikeFloat64 = ...,
+        out: None = ...,
+    ) -> float: ...
+    @overload
+    def standard_normal(  # type: ignore[misc]
+        self,
+        size: _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def standard_normal(  # type: ignore[misc]
+        self,
+        *,
+        out: ndarray[Any, dtype[float64]] = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def standard_normal(  # type: ignore[misc]
+        self,
+        size: _ShapeLike = ...,
+        dtype: _DTypeLikeFloat32 = ...,
+        out: None | ndarray[Any, dtype[float32]] = ...,
+    ) -> ndarray[Any, dtype[float32]]: ...
+    @overload
+    def standard_normal(  # type: ignore[misc]
+        self,
+        size: _ShapeLike = ...,
+        dtype: _DTypeLikeFloat64 = ...,
+        out: None | ndarray[Any, dtype[float64]] = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def permutation(self, x: int, axis: int = ...) -> ndarray[Any, dtype[int64]]: ...
+    @overload
+    def permutation(self, x: ArrayLike, axis: int = ...) -> ndarray[Any, Any]: ...
+    @overload
+    def standard_exponential(  # type: ignore[misc]
+        self,
+        size: None = ...,
+        dtype: _DTypeLikeFloat32 | _DTypeLikeFloat64 = ...,
+        method: Literal["zig", "inv"] = ...,
+        out: None = ...,
+    ) -> float: ...
+    @overload
+    def standard_exponential(
+        self,
+        size: _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def standard_exponential(
+        self,
+        *,
+        out: ndarray[Any, dtype[float64]] = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def standard_exponential(
+        self,
+        size: _ShapeLike = ...,
+        *,
+        method: Literal["zig", "inv"] = ...,
+        out: None | ndarray[Any, dtype[float64]] = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def standard_exponential(
+        self,
+        size: _ShapeLike = ...,
+        dtype: _DTypeLikeFloat32 = ...,
+        method: Literal["zig", "inv"] = ...,
+        out: None | ndarray[Any, dtype[float32]] = ...,
+    ) -> ndarray[Any, dtype[float32]]: ...
+    @overload
+    def standard_exponential(
+        self,
+        size: _ShapeLike = ...,
+        dtype: _DTypeLikeFloat64 = ...,
+        method: Literal["zig", "inv"] = ...,
+        out: None | ndarray[Any, dtype[float64]] = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def random(  # type: ignore[misc]
+        self,
+        size: None = ...,
+        dtype: _DTypeLikeFloat32 | _DTypeLikeFloat64 = ...,
+        out: None = ...,
+    ) -> float: ...
+    @overload
+    def random(
+        self,
+        *,
+        out: ndarray[Any, dtype[float64]] = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def random(
+        self,
+        size: _ShapeLike = ...,
+        *,
+        out: None | ndarray[Any, dtype[float64]] = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def random(
+        self,
+        size: _ShapeLike = ...,
+        dtype: _DTypeLikeFloat32 = ...,
+        out: None | ndarray[Any, dtype[float32]] = ...,
+    ) -> ndarray[Any, dtype[float32]]: ...
+    @overload
+    def random(
+        self,
+        size: _ShapeLike = ...,
+        dtype: _DTypeLikeFloat64 = ...,
+        out: None | ndarray[Any, dtype[float64]] = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def beta(
+        self,
+        a: _FloatLike_co,
+        b: _FloatLike_co,
+        size: None = ...,
+    ) -> float: ...  # type: ignore[misc]
+    @overload
+    def beta(
+        self, a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def exponential(self, scale: _FloatLike_co = ..., size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def exponential(
+        self, scale: _ArrayLikeFloat_co = ..., size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: int,
+        high: None | int = ...,
+    ) -> int: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: int,
+        high: None | int = ...,
+        size: None = ...,
+        dtype: _DTypeLikeBool = ...,
+        endpoint: bool = ...,
+    ) -> bool: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: int,
+        high: None | int = ...,
+        size: None = ...,
+        dtype: _DTypeLikeInt | _DTypeLikeUInt = ...,
+        endpoint: bool = ...,
+    ) -> int: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[int64]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+        dtype: _DTypeLikeBool = ...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[bool_]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+        dtype: dtype[int8] | type[int8] | _Int8Codes | _SupportsDType[dtype[int8]] = ...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[int8]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+        dtype: dtype[int16] | type[int16] | _Int16Codes | _SupportsDType[dtype[int16]] = ...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[int16]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+        dtype: dtype[int32] | type[int32] | _Int32Codes | _SupportsDType[dtype[int32]] = ...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[int32]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+        dtype: None | dtype[int64] | type[int64] | _Int64Codes | _SupportsDType[dtype[int64]] = ...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[int64]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+        dtype: dtype[uint8] | type[uint8] | _UInt8Codes | _SupportsDType[dtype[uint8]] = ...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[uint8]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+        dtype: dtype[uint16] | type[uint16] | _UInt16Codes | _SupportsDType[dtype[uint16]] = ...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[uint16]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+        dtype: dtype[uint32] | type[uint32] | _UInt32Codes | _SupportsDType[dtype[uint32]] = ...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[uint32]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+        dtype: dtype[uint64] | type[uint64] | _UInt64Codes | _SupportsDType[dtype[uint64]] = ...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[uint64]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+        dtype: dtype[int_] | type[int] | type[int_] | _IntCodes | _SupportsDType[dtype[int_]] = ...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[int_]]: ...
+    @overload
+    def integers(  # type: ignore[misc]
+        self,
+        low: _ArrayLikeInt_co,
+        high: None | _ArrayLikeInt_co = ...,
+        size: None | _ShapeLike = ...,
+        dtype: dtype[uint] | type[uint] | _UIntCodes | _SupportsDType[dtype[uint]] = ...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[uint]]: ...
+    # TODO: Use a TypeVar _T here to get away from Any output?  Should be int->ndarray[Any,dtype[int64]], ArrayLike[_T] -> _T | ndarray[Any,Any]
+    @overload
+    def choice(
+        self,
+        a: int,
+        size: None = ...,
+        replace: bool = ...,
+        p: None | _ArrayLikeFloat_co = ...,
+        axis: int = ...,
+        shuffle: bool = ...,
+    ) -> int: ...
+    @overload
+    def choice(
+        self,
+        a: int,
+        size: _ShapeLike = ...,
+        replace: bool = ...,
+        p: None | _ArrayLikeFloat_co = ...,
+        axis: int = ...,
+        shuffle: bool = ...,
+    ) -> ndarray[Any, dtype[int64]]: ...
+    @overload
+    def choice(
+        self,
+        a: ArrayLike,
+        size: None = ...,
+        replace: bool = ...,
+        p: None | _ArrayLikeFloat_co = ...,
+        axis: int = ...,
+        shuffle: bool = ...,
+    ) -> Any: ...
+    @overload
+    def choice(
+        self,
+        a: ArrayLike,
+        size: _ShapeLike = ...,
+        replace: bool = ...,
+        p: None | _ArrayLikeFloat_co = ...,
+        axis: int = ...,
+        shuffle: bool = ...,
+    ) -> ndarray[Any, Any]: ...
+    @overload
+    def uniform(
+        self,
+        low: _FloatLike_co = ...,
+        high: _FloatLike_co = ...,
+        size: None = ...,
+    ) -> float: ...  # type: ignore[misc]
+    @overload
+    def uniform(
+        self,
+        low: _ArrayLikeFloat_co = ...,
+        high: _ArrayLikeFloat_co = ...,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def normal(
+        self,
+        loc: _FloatLike_co = ...,
+        scale: _FloatLike_co = ...,
+        size: None = ...,
+    ) -> float: ...  # type: ignore[misc]
+    @overload
+    def normal(
+        self,
+        loc: _ArrayLikeFloat_co = ...,
+        scale: _ArrayLikeFloat_co = ...,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def standard_gamma(  # type: ignore[misc]
+        self,
+        shape: _FloatLike_co,
+        size: None = ...,
+        dtype: _DTypeLikeFloat32 | _DTypeLikeFloat64 = ...,
+        out: None = ...,
+    ) -> float: ...
+    @overload
+    def standard_gamma(
+        self,
+        shape: _ArrayLikeFloat_co,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def standard_gamma(
+        self,
+        shape: _ArrayLikeFloat_co,
+        *,
+        out: ndarray[Any, dtype[float64]] = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def standard_gamma(
+        self,
+        shape: _ArrayLikeFloat_co,
+        size: None | _ShapeLike = ...,
+        dtype: _DTypeLikeFloat32 = ...,
+        out: None | ndarray[Any, dtype[float32]] = ...,
+    ) -> ndarray[Any, dtype[float32]]: ...
+    @overload
+    def standard_gamma(
+        self,
+        shape: _ArrayLikeFloat_co,
+        size: None | _ShapeLike = ...,
+        dtype: _DTypeLikeFloat64 = ...,
+        out: None | ndarray[Any, dtype[float64]] = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def gamma(self, shape: _FloatLike_co, scale: _FloatLike_co = ..., size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def gamma(
+        self,
+        shape: _ArrayLikeFloat_co,
+        scale: _ArrayLikeFloat_co = ...,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def f(self, dfnum: _FloatLike_co, dfden: _FloatLike_co, size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def f(
+        self, dfnum: _ArrayLikeFloat_co, dfden: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def noncentral_f(self, dfnum: _FloatLike_co, dfden: _FloatLike_co, nonc: _FloatLike_co, size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def noncentral_f(
+        self,
+        dfnum: _ArrayLikeFloat_co,
+        dfden: _ArrayLikeFloat_co,
+        nonc: _ArrayLikeFloat_co,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def chisquare(self, df: _FloatLike_co, size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def chisquare(
+        self, df: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def noncentral_chisquare(self, df: _FloatLike_co, nonc: _FloatLike_co, size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def noncentral_chisquare(
+        self, df: _ArrayLikeFloat_co, nonc: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def standard_t(self, df: _FloatLike_co, size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def standard_t(
+        self, df: _ArrayLikeFloat_co, size: None = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def standard_t(
+        self, df: _ArrayLikeFloat_co, size: _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def vonmises(self, mu: _FloatLike_co, kappa: _FloatLike_co, size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def vonmises(
+        self, mu: _ArrayLikeFloat_co, kappa: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def pareto(self, a: _FloatLike_co, size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def pareto(
+        self, a: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def weibull(self, a: _FloatLike_co, size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def weibull(
+        self, a: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def power(self, a: _FloatLike_co, size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def power(
+        self, a: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def standard_cauchy(self, size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def standard_cauchy(self, size: _ShapeLike = ...) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def laplace(
+        self,
+        loc: _FloatLike_co = ...,
+        scale: _FloatLike_co = ...,
+        size: None = ...,
+    ) -> float: ...  # type: ignore[misc]
+    @overload
+    def laplace(
+        self,
+        loc: _ArrayLikeFloat_co = ...,
+        scale: _ArrayLikeFloat_co = ...,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def gumbel(
+        self,
+        loc: _FloatLike_co = ...,
+        scale: _FloatLike_co = ...,
+        size: None = ...,
+    ) -> float: ...  # type: ignore[misc]
+    @overload
+    def gumbel(
+        self,
+        loc: _ArrayLikeFloat_co = ...,
+        scale: _ArrayLikeFloat_co = ...,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def logistic(
+        self,
+        loc: _FloatLike_co = ...,
+        scale: _FloatLike_co = ...,
+        size: None = ...,
+    ) -> float: ...  # type: ignore[misc]
+    @overload
+    def logistic(
+        self,
+        loc: _ArrayLikeFloat_co = ...,
+        scale: _ArrayLikeFloat_co = ...,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def lognormal(
+        self,
+        mean: _FloatLike_co = ...,
+        sigma: _FloatLike_co = ...,
+        size: None = ...,
+    ) -> float: ...  # type: ignore[misc]
+    @overload
+    def lognormal(
+        self,
+        mean: _ArrayLikeFloat_co = ...,
+        sigma: _ArrayLikeFloat_co = ...,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def rayleigh(self, scale: _FloatLike_co = ..., size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def rayleigh(
+        self, scale: _ArrayLikeFloat_co = ..., size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def wald(self, mean: _FloatLike_co, scale: _FloatLike_co, size: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def wald(
+        self, mean: _ArrayLikeFloat_co, scale: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def triangular(
+        self,
+        left: _FloatLike_co,
+        mode: _FloatLike_co,
+        right: _FloatLike_co,
+        size: None = ...,
+    ) -> float: ...  # type: ignore[misc]
+    @overload
+    def triangular(
+        self,
+        left: _ArrayLikeFloat_co,
+        mode: _ArrayLikeFloat_co,
+        right: _ArrayLikeFloat_co,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    @overload
+    def binomial(self, n: int, p: _FloatLike_co, size: None = ...) -> int: ...  # type: ignore[misc]
+    @overload
+    def binomial(
+        self, n: _ArrayLikeInt_co, p: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[int64]]: ...
+    @overload
+    def negative_binomial(self, n: _FloatLike_co, p: _FloatLike_co, size: None = ...) -> int: ...  # type: ignore[misc]
+    @overload
+    def negative_binomial(
+        self, n: _ArrayLikeFloat_co, p: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[int64]]: ...
+    @overload
+    def poisson(self, lam: _FloatLike_co = ..., size: None = ...) -> int: ...  # type: ignore[misc]
+    @overload
+    def poisson(
+        self, lam: _ArrayLikeFloat_co = ..., size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[int64]]: ...
+    @overload
+    def zipf(self, a: _FloatLike_co, size: None = ...) -> int: ...  # type: ignore[misc]
+    @overload
+    def zipf(
+        self, a: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[int64]]: ...
+    @overload
+    def geometric(self, p: _FloatLike_co, size: None = ...) -> int: ...  # type: ignore[misc]
+    @overload
+    def geometric(
+        self, p: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[int64]]: ...
+    @overload
+    def hypergeometric(self, ngood: int, nbad: int, nsample: int, size: None = ...) -> int: ...  # type: ignore[misc]
+    @overload
+    def hypergeometric(
+        self,
+        ngood: _ArrayLikeInt_co,
+        nbad: _ArrayLikeInt_co,
+        nsample: _ArrayLikeInt_co,
+        size: None | _ShapeLike = ...,
+    ) -> ndarray[Any, dtype[int64]]: ...
+    @overload
+    def logseries(self, p: _FloatLike_co, size: None = ...) -> int: ...  # type: ignore[misc]
+    @overload
+    def logseries(
+        self, p: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[int64]]: ...
+    def multivariate_normal(
+        self,
+        mean: _ArrayLikeFloat_co,
+        cov: _ArrayLikeFloat_co,
+        size: None | _ShapeLike = ...,
+        check_valid: Literal["warn", "raise", "ignore"] = ...,
+        tol: float = ...,
+        *,
+        method: Literal["svd", "eigh", "cholesky"] = ...,
+    ) -> ndarray[Any, dtype[float64]]: ...
+    def multinomial(
+        self, n: _ArrayLikeInt_co,
+            pvals: _ArrayLikeFloat_co,
+            size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[int64]]: ...
+    def multivariate_hypergeometric(
+        self,
+        colors: _ArrayLikeInt_co,
+        nsample: int,
+        size: None | _ShapeLike = ...,
+        method: Literal["marginals", "count"] = ...,
+    ) -> ndarray[Any, dtype[int64]]: ...
+    def dirichlet(
+        self, alpha: _ArrayLikeFloat_co, size: None | _ShapeLike = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    def permuted(
+        self, x: ArrayLike, *, axis: None | int = ..., out: None | ndarray[Any, Any] = ...
+    ) -> ndarray[Any, Any]: ...
+    def shuffle(self, x: ArrayLike, axis: int = ...) -> None: ...
+
+def default_rng(
+    seed: None | _ArrayLikeInt_co | SeedSequence | BitGenerator | Generator = ...
+) -> Generator: ...

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/random/_pcg64.pyi

@@ -0,0 +1,42 @@
+from typing import TypedDict
+
+from numpy.random.bit_generator import BitGenerator, SeedSequence
+from numpy._typing import _ArrayLikeInt_co
+
+class _PCG64Internal(TypedDict):
+    state: int
+    inc: int
+
+class _PCG64State(TypedDict):
+    bit_generator: str
+    state: _PCG64Internal
+    has_uint32: int
+    uinteger: int
+
+class PCG64(BitGenerator):
+    def __init__(self, seed: None | _ArrayLikeInt_co | SeedSequence = ...) -> None: ...
+    def jumped(self, jumps: int = ...) -> PCG64: ...
+    @property
+    def state(
+        self,
+    ) -> _PCG64State: ...
+    @state.setter
+    def state(
+        self,
+        value: _PCG64State,
+    ) -> None: ...
+    def advance(self, delta: int) -> PCG64: ...
+
+class PCG64DXSM(BitGenerator):
+    def __init__(self, seed: None | _ArrayLikeInt_co | SeedSequence = ...) -> None: ...
+    def jumped(self, jumps: int = ...) -> PCG64DXSM: ...
+    @property
+    def state(
+        self,
+    ) -> _PCG64State: ...
+    @state.setter
+    def state(
+        self,
+        value: _PCG64State,
+    ) -> None: ...
+    def advance(self, delta: int) -> PCG64DXSM: ...

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/random/_pickle.py

@@ -0,0 +1,80 @@
+from .mtrand import RandomState
+from ._philox import Philox
+from ._pcg64 import PCG64, PCG64DXSM
+from ._sfc64 import SFC64
+
+from ._generator import Generator
+from ._mt19937 import MT19937
+
+BitGenerators = {'MT19937': MT19937,
+                 'PCG64': PCG64,
+                 'PCG64DXSM': PCG64DXSM,
+                 'Philox': Philox,
+                 'SFC64': SFC64,
+                 }
+
+
+def __bit_generator_ctor(bit_generator_name='MT19937'):
+    """
+    Pickling helper function that returns a bit generator object
+
+    Parameters
+    ----------
+    bit_generator_name : str
+        String containing the name of the BitGenerator
+
+    Returns
+    -------
+    bit_generator : BitGenerator
+        BitGenerator instance
+    """
+    if bit_generator_name in BitGenerators:
+        bit_generator = BitGenerators[bit_generator_name]
+    else:
+        raise ValueError(str(bit_generator_name) + ' is not a known '
+                                                   'BitGenerator module.')
+
+    return bit_generator()
+
+
+def __generator_ctor(bit_generator_name="MT19937",
+                     bit_generator_ctor=__bit_generator_ctor):
+    """
+    Pickling helper function that returns a Generator object
+
+    Parameters
+    ----------
+    bit_generator_name : str
+        String containing the core BitGenerator's name
+    bit_generator_ctor : callable, optional
+        Callable function that takes bit_generator_name as its only argument
+        and returns an instantized bit generator.
+
+    Returns
+    -------
+    rg : Generator
+        Generator using the named core BitGenerator
+    """
+    return Generator(bit_generator_ctor(bit_generator_name))
+
+
+def __randomstate_ctor(bit_generator_name="MT19937",
+                       bit_generator_ctor=__bit_generator_ctor):
+    """
+    Pickling helper function that returns a legacy RandomState-like object
+
+    Parameters
+    ----------
+    bit_generator_name : str
+        String containing the core BitGenerator's name
+    bit_generator_ctor : callable, optional
+        Callable function that takes bit_generator_name as its only argument
+        and returns an instantized bit generator.
+
+    Returns
+    -------
+    rs : RandomState
+        Legacy RandomState using the named core BitGenerator
+    """
+
+    return RandomState(bit_generator_ctor(bit_generator_name))

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/random/_sfc64.pyi

@@ -0,0 +1,28 @@
+from typing import Any, TypedDict
+
+from numpy import dtype as dtype
+from numpy import ndarray as ndarray
+from numpy import uint64
+from numpy.random.bit_generator import BitGenerator, SeedSequence
+from numpy._typing import _ArrayLikeInt_co
+
+class _SFC64Internal(TypedDict):
+    state: ndarray[Any, dtype[uint64]]
+
+class _SFC64State(TypedDict):
+    bit_generator: str
+    state: _SFC64Internal
+    has_uint32: int
+    uinteger: int
+
+class SFC64(BitGenerator):
+    def __init__(self, seed: None | _ArrayLikeInt_co | SeedSequence = ...) -> None: ...
+    @property
+    def state(
+        self,
+    ) -> _SFC64State: ...
+    @state.setter
+    def state(
+        self,
+        value: _SFC64State,
+    ) -> None: ...

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/random/c_distributions.pxd

@@ -0,0 +1,120 @@
+#!python
+#cython: wraparound=False, nonecheck=False, boundscheck=False, cdivision=True, language_level=3
+from numpy cimport npy_intp
+
+from libc.stdint cimport (uint64_t, int32_t, int64_t)
+from numpy.random cimport bitgen_t
+
+cdef extern from "numpy/random/distributions.h":
+
+    struct s_binomial_t:
+        int has_binomial
+        double psave
+        int64_t nsave
+        double r
+        double q
+        double fm
+        int64_t m
+        double p1
+        double xm
+        double xl
+        double xr
+        double c
+        double laml
+        double lamr
+        double p2
+        double p3
+        double p4
+
+    ctypedef s_binomial_t binomial_t
+
+    float random_standard_uniform_f(bitgen_t *bitgen_state) nogil
+    double random_standard_uniform(bitgen_t *bitgen_state) nogil
+    void random_standard_uniform_fill(bitgen_t* bitgen_state, npy_intp cnt, double *out) nogil
+    void random_standard_uniform_fill_f(bitgen_t *bitgen_state, npy_intp cnt, float *out) nogil
+    
+    double random_standard_exponential(bitgen_t *bitgen_state) nogil
+    float random_standard_exponential_f(bitgen_t *bitgen_state) nogil
+    void random_standard_exponential_fill(bitgen_t *bitgen_state, npy_intp cnt, double *out) nogil
+    void random_standard_exponential_fill_f(bitgen_t *bitgen_state, npy_intp cnt, float *out) nogil
+    void random_standard_exponential_inv_fill(bitgen_t *bitgen_state, npy_intp cnt, double *out) nogil
+    void random_standard_exponential_inv_fill_f(bitgen_t *bitgen_state, npy_intp cnt, float *out) nogil
+    
+    double random_standard_normal(bitgen_t* bitgen_state) nogil
+    float random_standard_normal_f(bitgen_t *bitgen_state) nogil
+    void random_standard_normal_fill(bitgen_t *bitgen_state, npy_intp count, double *out) nogil
+    void random_standard_normal_fill_f(bitgen_t *bitgen_state, npy_intp count, float *out) nogil
+    double random_standard_gamma(bitgen_t *bitgen_state, double shape) nogil
+    float random_standard_gamma_f(bitgen_t *bitgen_state, float shape) nogil
+
+    float random_standard_uniform_f(bitgen_t *bitgen_state) nogil
+    void random_standard_uniform_fill_f(bitgen_t* bitgen_state, npy_intp cnt, float *out) nogil
+    float random_standard_normal_f(bitgen_t* bitgen_state) nogil
+    float random_standard_gamma_f(bitgen_t *bitgen_state, float shape) nogil
+
+    int64_t random_positive_int64(bitgen_t *bitgen_state) nogil
+    int32_t random_positive_int32(bitgen_t *bitgen_state) nogil
+    int64_t random_positive_int(bitgen_t *bitgen_state) nogil
+    uint64_t random_uint(bitgen_t *bitgen_state) nogil
+
+    double random_normal(bitgen_t *bitgen_state, double loc, double scale) nogil
+
+    double random_gamma(bitgen_t *bitgen_state, double shape, double scale) nogil
+    float random_gamma_f(bitgen_t *bitgen_state, float shape, float scale) nogil
+
+    double random_exponential(bitgen_t *bitgen_state, double scale) nogil
+    double random_uniform(bitgen_t *bitgen_state, double lower, double range) nogil
+    double random_beta(bitgen_t *bitgen_state, double a, double b) nogil
+    double random_chisquare(bitgen_t *bitgen_state, double df) nogil
+    double random_f(bitgen_t *bitgen_state, double dfnum, double dfden) nogil
+    double random_standard_cauchy(bitgen_t *bitgen_state) nogil
+    double random_pareto(bitgen_t *bitgen_state, double a) nogil
+    double random_weibull(bitgen_t *bitgen_state, double a) nogil
+    double random_power(bitgen_t *bitgen_state, double a) nogil
+    double random_laplace(bitgen_t *bitgen_state, double loc, double scale) nogil
+    double random_gumbel(bitgen_t *bitgen_state, double loc, double scale) nogil
+    double random_logistic(bitgen_t *bitgen_state, double loc, double scale) nogil
+    double random_lognormal(bitgen_t *bitgen_state, double mean, double sigma) nogil
+    double random_rayleigh(bitgen_t *bitgen_state, double mode) nogil
+    double random_standard_t(bitgen_t *bitgen_state, double df) nogil
+    double random_noncentral_chisquare(bitgen_t *bitgen_state, double df,
+                                       double nonc) nogil
+    double random_noncentral_f(bitgen_t *bitgen_state, double dfnum,
+                               double dfden, double nonc) nogil
+    double random_wald(bitgen_t *bitgen_state, double mean, double scale) nogil
+    double random_vonmises(bitgen_t *bitgen_state, double mu, double kappa) nogil
+    double random_triangular(bitgen_t *bitgen_state, double left, double mode,
+                             double right) nogil
+
+    int64_t random_poisson(bitgen_t *bitgen_state, double lam) nogil
+    int64_t random_negative_binomial(bitgen_t *bitgen_state, double n, double p) nogil
+    int64_t random_binomial(bitgen_t *bitgen_state, double p, int64_t n, binomial_t *binomial) nogil
+    int64_t random_logseries(bitgen_t *bitgen_state, double p) nogil
+    int64_t random_geometric_search(bitgen_t *bitgen_state, double p) nogil
+    int64_t random_geometric_inversion(bitgen_t *bitgen_state, double p) nogil
+    int64_t random_geometric(bitgen_t *bitgen_state, double p) nogil
+    int64_t random_zipf(bitgen_t *bitgen_state, double a) nogil
+    int64_t random_hypergeometric(bitgen_t *bitgen_state, int64_t good, int64_t bad,
+                                    int64_t sample) nogil
+
+    uint64_t random_interval(bitgen_t *bitgen_state, uint64_t max) nogil
+
+    # Generate random uint64 numbers in closed interval [off, off + rng].
+    uint64_t random_bounded_uint64(bitgen_t *bitgen_state,
+                                   uint64_t off, uint64_t rng,
+                                   uint64_t mask, bint use_masked) nogil
+
+    void random_multinomial(bitgen_t *bitgen_state, int64_t n, int64_t *mnix,
+                            double *pix, npy_intp d, binomial_t *binomial) nogil
+
+    int random_multivariate_hypergeometric_count(bitgen_t *bitgen_state,
+                          int64_t total,
+                          size_t num_colors, int64_t *colors,
+                          int64_t nsample,
+                          size_t num_variates, int64_t *variates) nogil
+    void random_multivariate_hypergeometric_marginals(bitgen_t *bitgen_state,
+                               int64_t total,
+                               size_t num_colors, int64_t *colors,
+                               int64_t nsample,
+                               size_t num_variates, int64_t *variates) nogil
+

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/fsmt/__init__.py

@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_fsmt import *
+    from .modeling_fsmt import *
+    from .tokenization_fsmt import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/fsmt/modeling_fsmt.py

@@ -0,0 +1,1136 @@
+# Copyright 2020 The Facebook AI Research Team Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Original implementation: https://github.com/pytorch/fairseq/tree/master/examples/wmt19
+# Authors:
+# - @alexeib Alexei Baevski
+# - @edunov Sergey Edunov
+# - @michaelauli Michael Auli
+# - @myleott Myle Ott
+# - @nng555 Nathan Ng
+# - David Grangier
+# - Kyra Yee
+#
+# Paper: Facebook FAIR's WMT19 News Translation Task Submission https://huggingface.co/papers/1907.06616
+#
+"""PyTorch Fairseq model, ported from https://github.com/pytorch/fairseq/tree/master/examples/wmt19"""
+
+import math
+from typing import Any
+
+import torch
+from torch import Tensor, nn
+from torch.nn import CrossEntropyLoss, LayerNorm
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...modeling_outputs import (
+    BaseModelOutput,
+    BaseModelOutputWithPastAndCrossAttentions,
+    Seq2SeqLMOutput,
+    Seq2SeqModelOutput,
+)
+from ...modeling_utils import PreTrainedModel
+from ...utils import auto_docstring, logging
+from .configuration_fsmt import FSMTConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# See all FSMT models at https://huggingface.co/models?filter=fsmt
+
+# Porting notes:
+# this one is modeled after BartModel*
+#
+# Currently only translation (fairseq also has weights for LM)
+#
+# fairseq provides weights for ru-en, en-ru and de-en, en-de pairs. All have been ported.
+# - ru-en, en-ru use asymmetric vocab
+# - de-en, en-de use a merged single vocab (but the code works as if they are separate)
+#
+# Differences with Bart:
+# - not using bos token
+# - 2 separate vocabs (src and target)
+# - embed weights aren't tied
+# - uses a model Ensemble (but that part isn't ported/implemented yet) - so we
+#   aren't getting as good of a BLEU score
+# - uses a projection layer at the end of the decoder
+# - doesn't use final_logits_bias
+# - beam search: stops as soon as num_beams == len(hypos) (whereas transformers
+#   is not satisfied there and will continue searching until the next cycles
+#   aren't promising something better), comparing BLEU scores - the transformers
+#   algorithm is slightly superior, therefore using the latter. But if you want
+#   to match fairseq outputs, you need to pass ``early_stopping=True`` to ``generate()``.
+#
+# SinusoidalPositionalEmbedding is slightly different from Bart's - generates
+# different embeddings. This implementation is copied verbatim from fairseq with
+# some small changes to make it work here.
+#
+# Other changes:
+#  - doesn't support use_cache as Bart's version does
+#
+#
+# FSMTConfig changes with BartConfig
+#
+#    Differences with BART:
+#    - src/tgt vocabs aren't shared
+#    - token embeddings aren't shared
+#    - needs a language pair
+#    - scale_embedding are True
+#
+#    some unused args were removed too
+#
+#
+# TODO:
+# - port model ensemble (fs uses 4 model checkpoints)
+# - solve beam search discrepancies
+# docstyle-ignore
+
+"""
+
+Here is how to compare BLEU scores against fairseq implementation:
+(don't forget to install sacrebleu: `pip install sacrebleu`)
+
+# en-ru
+
+export PAIR=en-ru
+export DATA_DIR=data/$PAIR
+export SAVE_DIR=data/$PAIR
+export BS=8
+export NUM_BEAMS=50
+mkdir -p $DATA_DIR
+sacrebleu -t wmt19 -l $PAIR --echo src > $DATA_DIR/val.source
+sacrebleu -t wmt19 -l $PAIR --echo ref > $DATA_DIR/val.target
+echo $PAIR
+PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval.py facebook/wmt19-$PAIR $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation --num_beams $NUM_BEAMS
+
+# (fairseq BLEU: 36.4 http://matrix.statmt.org/matrix/output/1914?score_id=37605)
+
+
+# ru-en
+
+export PAIR=ru-en
+export DATA_DIR=data/$PAIR
+export SAVE_DIR=data/$PAIR
+export BS=8
+export NUM_BEAMS=50
+mkdir -p $DATA_DIR
+sacrebleu -t wmt19 -l $PAIR --echo src > $DATA_DIR/val.source
+sacrebleu -t wmt19 -l $PAIR --echo ref > $DATA_DIR/val.target
+PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval.py facebook/wmt19-$PAIR $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation --num_beams $NUM_BEAMS
+
+
+# (fairseq BLEU: 41.3 http://matrix.statmt.org/matrix/output/1907?run_id=6937)
+
+
+# de-en
+
+export PAIR=de-en
+export DATA_DIR=data/$PAIR
+export SAVE_DIR=data/$PAIR
+export BS=8
+export NUM_BEAMS=50
+mkdir -p $DATA_DIR
+sacrebleu -t wmt19 -l $PAIR --echo src > $DATA_DIR/val.source
+sacrebleu -t wmt19 -l $PAIR --echo ref > $DATA_DIR/val.target
+echo $PAIR
+PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval.py facebook/wmt19-$PAIR $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation --num_beams $NUM_BEAMS
+
+# (fairseq BLEU: 42.3 http://matrix.statmt.org/matrix/output/1902?run_id=6750)
+
+
+
+# en-de
+
+export PAIR=en-de
+export DATA_DIR=data/$PAIR
+export SAVE_DIR=data/$PAIR
+export BS=8
+mkdir -p $DATA_DIR
+sacrebleu -t wmt19 -l $PAIR --echo src > $DATA_DIR/val.source
+sacrebleu -t wmt19 -l $PAIR --echo ref > $DATA_DIR/val.target
+echo $PAIR
+PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval.py facebook/wmt19-$PAIR $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation --num_beams $NUM_BEAMS
+
+# (fairseq BLEU: 43.1 http://matrix.statmt.org/matrix/output/1909?run_id=6862)
+
+"""
+
+
+def invert_mask(attention_mask):
+    """Turns 1->0, 0->1, False->True, True-> False"""
+    assert attention_mask.dim() == 2
+    return attention_mask.eq(0)
+
+
+def triu_onnx(x, diagonal=0):
+    l = x.shape[0]
+    arange = torch.arange(l, device=x.device)
+    mask = arange.expand(l, l)
+    arange = arange.unsqueeze(-1)
+    if diagonal:
+        arange = arange + diagonal
+    mask = mask >= arange
+    return x.masked_fill(mask == 0, 0)
+
+
+def _prepare_fsmt_decoder_inputs(
+    config,
+    input_ids,
+    decoder_input_ids=None,
+    decoder_padding_mask=None,
+    causal_mask_dtype=torch.float32,
+):
+    """
+    Prepare masks that ignore padding tokens in the decoder and a causal mask for the decoder if none are provided.
+    This mimics the default behavior in fairseq. To override it pass in masks. Note: this is not called during
+    generation
+    """
+    pad_token_id = config.pad_token_id
+    if decoder_input_ids is None:
+        decoder_input_ids = shift_tokens_right(input_ids, pad_token_id)
+    bsz, tgt_len = decoder_input_ids.size()
+    if decoder_padding_mask is None:
+        decoder_padding_mask = make_padding_mask(decoder_input_ids, pad_token_id)
+    else:
+        decoder_padding_mask = invert_mask(decoder_padding_mask)
+    causal_mask = triu_onnx(fill_with_neg_inf(torch.zeros(tgt_len, tgt_len, dtype=causal_mask_dtype)), 1).to(
+        device=decoder_input_ids.device
+    )
+    return decoder_input_ids, decoder_padding_mask, causal_mask
+
+
+@auto_docstring
+class PretrainedFSMTModel(PreTrainedModel):
+    config: FSMTConfig
+    base_model_prefix = "model"
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        std = self.config.init_std
+        if isinstance(module, nn.Linear):
+            init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                init.zeros_(module.bias)
+        elif isinstance(module, SinusoidalPositionalEmbedding):
+            weight = module.get_embedding(*module.weight.shape, module.padding_idx)
+            init.copy_(module.weight, weight)
+        elif isinstance(module, nn.Embedding):
+            init.normal_(module.weight, mean=0.0, std=std)
+            # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag
+            if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
+                init.zeros_(module.weight[module.padding_idx])
+
+    @property
+    def dummy_inputs(self):
+        pad_token = self.config.pad_token_id
+        input_ids = torch.tensor([[0, 6, 10, 4, 2], [0, 8, 12, 2, pad_token]], device=self.device)
+        dummy_inputs = {
+            "attention_mask": input_ids.ne(pad_token),
+            "input_ids": input_ids,
+        }
+        return dummy_inputs
+
+
+def _make_linear_from_emb(emb):
+    vocab_size, emb_size = emb.weight.shape
+    lin_layer = nn.Linear(vocab_size, emb_size, bias=False)
+    lin_layer.weight.data = emb.weight.data
+    return lin_layer
+
+
+# Helper Functions, mostly for making masks
+def _check_shapes(shape_1, shape2):
+    if shape_1 != shape2:
+        raise AssertionError(f"shape mismatch: {shape_1} != {shape2}")
+
+
+def shift_tokens_right(input_ids, pad_token_id):
+    """Shift input ids one token to the right, and wrap the last non pad token (usually <eos>)."""
+
+    # replace possible -100 values in labels by `pad_token_id`
+    input_ids.masked_fill_(input_ids == -100, pad_token_id)
+
+    prev_output_tokens = input_ids.clone()
+    index_of_eos = (input_ids.ne(pad_token_id).sum(dim=1) - 1).unsqueeze(-1)
+    prev_output_tokens[:, 0] = input_ids.gather(1, index_of_eos).squeeze()
+    prev_output_tokens[:, 1:] = input_ids[:, :-1]
+    return prev_output_tokens
+
+
+def make_padding_mask(input_ids, padding_idx=1):
+    """True for pad tokens"""
+    padding_mask = input_ids.eq(padding_idx)
+    if not padding_mask.any():
+        padding_mask = None
+    return padding_mask
+
+
+# Helper Modules
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, config: FSMTConfig):
+        super().__init__()
+        self.embed_dim = config.d_model
+        self.self_attn = Attention(self.embed_dim, config.encoder_attention_heads, dropout=config.attention_dropout)
+        self.self_attn_layer_norm = LayerNorm(self.embed_dim)
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.activation_function]
+        self.activation_dropout = config.activation_dropout
+        self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+        self.final_layer_norm = LayerNorm(self.embed_dim)
+
+    def forward(self, x, encoder_padding_mask, output_attentions=False):
+        """
+        Args:
+            x (`torch.Tensor`): input to the layer of shape *(seq_len, batch, embed_dim)*
+            encoder_padding_mask (`torch.ByteTensor`): binary ByteTensor of shape
+                *(batch, src_len)* where padding elements are indicated by `1`.
+            for t_tgt, t_src is excluded (or masked out), =0 means it is
+            included in attention
+
+        Returns:
+            encoded output of shape *(seq_len, batch, embed_dim)*
+        """
+        residual = x
+        x, attn_weights = self.self_attn(
+            query=x,
+            key=x,
+            key_padding_mask=encoder_padding_mask,
+            output_attentions=output_attentions,
+        )
+        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
+        x = residual + x
+        x = self.self_attn_layer_norm(x)
+
+        residual = x
+        x = self.activation_fn(self.fc1(x))
+        x = nn.functional.dropout(x, p=self.activation_dropout, training=self.training)
+        x = self.fc2(x)
+        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
+        x = residual + x
+        x = self.final_layer_norm(x)
+        return x, attn_weights
+
+
+class FSMTEncoder(nn.Module):
+    """
+    Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a [`EncoderLayer`].
+
+    Args:
+        config: FSMTConfig
+    """
+
+    def __init__(self, config: FSMTConfig):
+        super().__init__()
+        self.dropout = config.dropout
+        self.layerdrop = config.encoder_layerdrop
+        self.padding_idx = config.pad_token_id
+        self.embed_tokens = nn.Embedding(config.src_vocab_size, config.d_model, config.pad_token_id)
+        embed_dim = self.embed_tokens.embedding_dim
+        self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+        self.embed_positions = SinusoidalPositionalEmbedding(
+            config.max_position_embeddings + self.padding_idx + 1, embed_dim, self.padding_idx
+        )
+        self.layers = nn.ModuleList([EncoderLayer(config) for _ in range(config.encoder_layers)])  # type: list[EncoderLayer]
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+    ):
+        """
+        Args:
+            input_ids (`torch.LongTensor`): tokens in the source language of shape
+                *(batch, src_len)*
+            attention_mask (`torch.LongTensor`): indicating which indices are padding tokens
+            inputs_embeds (`torch.FloatTensor`):
+                embedding vectors of shape *(batch, src_len, embed_dim)*
+
+        Returns:
+            BaseModelOutput or Tuple comprised of:
+
+                - **x** (`torch.Tensor`): the last encoder layer's output of shape *(src_len, batch, embed_dim)*
+                - **encoder_states** (`Tuple(torch.FloatTensor)`): all intermediate hidden states of shape *(src_len,
+                  batch, embed_dim)*. Only populated if *output_hidden_states:* is True.
+                - **all_attentions** (`Tuple(torch.FloatTensor)`): Attention weights for each layer.
+                During training might not be of length n_layers because of layer dropout.
+        """
+        # check attention mask and invert
+        if attention_mask is not None:
+            attention_mask = invert_mask(attention_mask)
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
+            embed_pos = self.embed_positions(input_ids)
+        elif inputs_embeds is not None:
+            inputs_embeds = inputs_embeds * self.embed_scale
+
+            # We assume zeros hidden states correspond to padding tokens
+            # and create `position_ids` where inputs_embeds[:, :, 0] == 0
+            position_ids = inputs_embeds[:, :, 0].masked_fill(
+                inputs_embeds[:, :, 0].eq(0), self.embed_positions.padding_idx
+            )
+
+            embed_pos = self.embed_positions(position_ids)
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        x = inputs_embeds + embed_pos
+        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
+
+        # B x T x C -> T x B x C
+        x = x.transpose(0, 1)
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                x = x.transpose(0, 1)  # T x B x C -> B x T x C
+                encoder_states += (x,)
+                x = x.transpose(0, 1)  # B x T x C -> T x B x C
+            # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+            dropout_probability = torch.rand([])
+            if self.training and (dropout_probability < self.layerdrop):  # skip the layer
+                attn = None
+            else:
+                x, attn = encoder_layer(
+                    x,
+                    attention_mask,
+                    output_attentions=output_attentions,
+                )
+
+            if output_attentions:
+                all_attentions = all_attentions + (attn,)
+
+        # T x B x C -> B x T x C
+        x = x.transpose(0, 1)
+
+        if output_hidden_states:
+            encoder_states += (x,)
+
+        if not return_dict:
+            return tuple(v for v in [x, encoder_states, all_attentions] if v is not None)
+        return BaseModelOutput(last_hidden_state=x, hidden_states=encoder_states, attentions=all_attentions)
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: FSMTConfig, layer_idx=None):
+        super().__init__()
+        self.embed_dim = config.d_model
+
+        self.self_attn = Attention(
+            embed_dim=self.embed_dim,
+            num_heads=config.decoder_attention_heads,
+            dropout=config.attention_dropout,
+            layer_idx=layer_idx,
+        )
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.activation_function]
+        self.activation_dropout = config.activation_dropout
+
+        self.self_attn_layer_norm = LayerNorm(self.embed_dim)
+        self.encoder_attn = Attention(
+            self.embed_dim,
+            config.decoder_attention_heads,
+            dropout=config.attention_dropout,
+            encoder_decoder_attention=True,
+            layer_idx=layer_idx,
+        )
+        self.encoder_attn_layer_norm = LayerNorm(self.embed_dim)
+        self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
+        self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
+        self.final_layer_norm = LayerNorm(self.embed_dim)
+
+    def forward(
+        self,
+        x,
+        encoder_hidden_states,
+        encoder_attn_mask=None,
+        layer_state=None,
+        causal_mask=None,
+        decoder_padding_mask=None,
+        output_attentions=False,
+        **kwargs,
+    ):
+        residual = x
+
+        # Self Attention
+        x, self_attn_weights = self.self_attn(
+            query=x,
+            key=x,
+            layer_state=layer_state,  # adds keys to layer state
+            key_padding_mask=decoder_padding_mask,
+            attn_mask=causal_mask,
+            output_attentions=output_attentions,
+        )
+        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
+        x = residual + x
+        x = self.self_attn_layer_norm(x)
+
+        # Cross attention
+        residual = x
+        assert self.encoder_attn.cache_key != self.self_attn.cache_key
+        x, cross_attn_weights = self.encoder_attn(
+            query=x,
+            key=encoder_hidden_states,
+            key_padding_mask=encoder_attn_mask,
+            layer_state=layer_state,  # mutates layer state
+            output_attentions=output_attentions,
+        )
+        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
+        x = residual + x
+        x = self.encoder_attn_layer_norm(x)
+
+        # Fully Connected
+        residual = x
+        x = self.activation_fn(self.fc1(x))
+        x = nn.functional.dropout(x, p=self.activation_dropout, training=self.training)
+        x = self.fc2(x)
+        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
+        x = residual + x
+        x = self.final_layer_norm(x)
+        return (
+            x,
+            self_attn_weights,
+            cross_attn_weights,
+        )
+
+
+class FSMTDecoder(nn.Module):
+    """
+    Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`DecoderLayer`]
+
+    Args:
+        config: FSMTConfig
+        embed_tokens (nn.Embedding): output embedding
+    """
+
+    def __init__(self, config: FSMTConfig):
+        super().__init__()
+        self.dropout = config.dropout
+        self.layerdrop = config.decoder_layerdrop
+        self.padding_idx = config.pad_token_id
+        self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+        self.embed_tokens = nn.Embedding(config.tgt_vocab_size, config.d_model, self.padding_idx)
+        embed_dim = self.embed_tokens.embedding_dim
+        self.embed_positions = SinusoidalPositionalEmbedding(
+            config.max_position_embeddings + self.padding_idx + 1, embed_dim, self.padding_idx
+        )
+        self.layers = nn.ModuleList([DecoderLayer(config, layer_idx=i) for i in range(config.decoder_layers)])  # type: list[DecoderLayer]
+        self.output_projection = nn.Linear(config.d_model, config.tgt_vocab_size, bias=False)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_padding_mask: torch.Tensor,
+        decoder_padding_mask: torch.Tensor,
+        decoder_causal_mask: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        output_attentions: bool | None = False,
+        output_hidden_states: bool | None = False,
+        return_dict: bool | None = True,
+        **kwargs,
+    ):
+        """
+        Includes several features from "Jointly Learning to Align and Translate with Transformer Models" (Garg et al.,
+        EMNLP 2019).
+
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch, tgt_len)`):
+                previous decoder outputs for teacher forcing
+            encoder_hidden_states: output from the encoder, used for
+                encoder-side attention
+            encoder_padding_mask: for ignoring pad tokens
+            past_key_values (dict or None): dictionary used for storing state during generation
+
+        Returns:
+            BaseModelOutputWithPast or tuple:
+
+                - the decoder's features of shape *(batch, tgt_len, embed_dim)*
+                - the cache
+                - hidden states
+                - attentions
+        """
+        # check attention mask and invert
+        if encoder_padding_mask is not None:
+            encoder_padding_mask = invert_mask(encoder_padding_mask)
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+        elif input_ids is not None:
+            # embed positions
+            positions = self.embed_positions(input_ids)
+            if use_cache:
+                input_ids = input_ids[:, -1:]
+                positions = positions[:, -1:]  # happens after we embed them
+            x = self.embed_tokens(input_ids) * self.embed_scale
+        elif inputs_embeds is not None:
+            # We assume zeros hidden states correspond to padding tokens
+            # and create `position_ids` where inputs_embeds[:, :, 0] == 0
+            position_ids = inputs_embeds[:, :, 0].masked_fill(
+                inputs_embeds[:, :, 0].eq(0), self.embed_positions.padding_idx
+            )
+            positions = self.embed_positions(position_ids)
+            x = inputs_embeds * self.embed_scale
+        else:
+            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+        x += positions
+        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
+
+        # Convert to FSMT output format: (BS, seq_len, model_dim) -> (seq_len, BS, model_dim)
+        x = x.transpose(0, 1)
+        encoder_hidden_states = encoder_hidden_states.transpose(0, 1)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        all_cross_attns = () if output_attentions else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+            if output_hidden_states:
+                x = x.transpose(0, 1)
+                all_hidden_states += (x,)
+                x = x.transpose(0, 1)
+            if self.training:
+                dropout_probability = torch.rand([])
+                if dropout_probability < self.layerdrop:
+                    continue
+
+            x, layer_self_attn, layer_cross_attn = decoder_layer(
+                x,
+                encoder_hidden_states,
+                encoder_attn_mask=encoder_padding_mask,
+                decoder_padding_mask=decoder_padding_mask,
+                layer_state=past_key_values,
+                causal_mask=decoder_causal_mask,
+                output_attentions=output_attentions,
+            )
+
+            if output_attentions:
+                all_self_attns += (layer_self_attn,)
+                all_cross_attns += (layer_cross_attn,)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            x = x.transpose(0, 1)
+            all_hidden_states += (x,)
+            x = x.transpose(0, 1)
+
+        # Convert to standard output format: (seq_len, BS, model_dim) -> (BS, seq_len, model_dim)
+        x = x.transpose(0, 1)
+        encoder_hidden_states = encoder_hidden_states.transpose(0, 1)
+
+        x = self.output_projection(x)
+
+        if not return_dict:
+            return tuple(
+                v for v in [x, past_key_values, all_hidden_states, all_self_attns, all_cross_attns] if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=x,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+            cross_attentions=all_cross_attns,
+        )
+
+
+def _reorder_buffer(attn_cache, new_order):
+    for k, input_buffer_k in attn_cache.items():
+        if input_buffer_k is not None:
+            attn_cache[k] = input_buffer_k.index_select(0, new_order)
+    return attn_cache
+
+
+class Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        bias=True,
+        encoder_decoder_attention=False,  # otherwise self_attention
+        layer_idx=None,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+        self.scaling = self.head_dim**-0.5
+        self.layer_idx = layer_idx
+
+        self.encoder_decoder_attention = encoder_decoder_attention
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.cache_key = "encoder_decoder" if self.encoder_decoder_attention else "self"
+
+    def forward(
+        self,
+        query,
+        key: Tensor | None,
+        key_padding_mask: Tensor | None = None,
+        layer_state: Cache | None = None,
+        attn_mask: Tensor | None = None,
+        output_attentions: bool | None = False,
+        **kwargs,
+    ) -> tuple[Tensor, Tensor | None]:
+        """Input shape: Time(SeqLen) x Batch x Channel"""
+        tgt_len, bsz, embed_dim = query.size()
+        assert embed_dim == self.embed_dim
+        assert list(query.size()) == [tgt_len, bsz, embed_dim]
+
+        if layer_state is not None:
+            if isinstance(layer_state, EncoderDecoderCache):
+                is_updated = layer_state.is_updated.get(self.layer_idx)
+                if self.encoder_decoder_attention:
+                    # after the first generated id, we can subsequently re-use all key/value_states from cache
+                    curr_past_key_values = layer_state.cross_attention_cache
+                else:
+                    curr_past_key_values = layer_state.self_attention_cache
+            else:
+                curr_past_key_values = layer_state
+
+        # NOTE: FSMT has format (seq_len, BS, model_dim) for inputs
+        current_states = key if self.encoder_decoder_attention else query
+        if self.encoder_decoder_attention and layer_state is not None and is_updated:
+            # reuse k,v, cross_attentions
+            key_states = curr_past_key_values.layers[self.layer_idx].keys
+            value_states = curr_past_key_values.layers[self.layer_idx].values
+        else:
+            key_states = self.k_proj(current_states)
+            value_states = self.v_proj(current_states)
+            key_states = key_states.view(-1, bsz, self.num_heads, self.head_dim).permute(1, 2, 0, 3)
+            value_states = value_states.view(-1, bsz, self.num_heads, self.head_dim).permute(1, 2, 0, 3)
+
+            if layer_state is not None:
+                # save all key/value_states to cache to be re-used for fast auto-regressive generation
+                key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)
+                # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+                if self.encoder_decoder_attention:
+                    layer_state.is_updated[self.layer_idx] = True
+
+        query_states = self.q_proj(query) * self.scaling
+
+        # Reshape back to 3D tensors for `bmm`
+        query_states = query_states.view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        key_states = key_states.reshape(bsz * self.num_heads, -1, self.head_dim)
+        value_states = value_states.reshape(bsz * self.num_heads, -1, self.head_dim)
+
+        assert key_states is not None
+        src_len = key_states.size(1)
+        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
+        assert attn_weights.size() == (bsz * self.num_heads, tgt_len, src_len)
+
+        if attn_mask is not None:
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attn_mask
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        # This is part of a workaround to get around fork/join parallelism not supporting Optional types.
+        if key_padding_mask is not None and key_padding_mask.dim() == 0:
+            key_padding_mask = None
+        assert key_padding_mask is None or key_padding_mask.size()[:2] == (
+            bsz,
+            src_len,
+        )
+
+        if key_padding_mask is not None:  # don't attend to padding symbols
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            reshaped = key_padding_mask.unsqueeze(1).unsqueeze(2)
+            attn_weights = attn_weights.masked_fill(reshaped, torch.finfo(attn_weights.dtype).min)
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+
+        if output_attentions:
+            # make sure that attn_weights are included in graph
+            attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
+        else:
+            attn_weights_reshaped = None
+
+        attn_probs = nn.functional.dropout(
+            attn_weights,
+            p=self.dropout,
+            training=self.training,
+        )
+
+        assert value_states is not None
+        attn_output = torch.bmm(attn_probs, value_states)
+        assert attn_output.size() == (bsz * self.num_heads, tgt_len, self.head_dim)
+        attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights_reshaped
+
+
+def fill_with_neg_inf(t):
+    """FP16-compatible function that fills a input_ids with -inf."""
+    return t.float().fill_(torch.finfo(t.dtype).min).type_as(t)
+
+
+# Public API
+def _get_shape(t):
+    return getattr(t, "shape", None)
+
+
+@auto_docstring
+class FSMTModel(PretrainedFSMTModel):
+    _tied_weights_keys = {
+        "encoder.embed_tokens.weight": "decoder.embed_tokens.weight",
+        "decoder.output_projection.weight": "decoder.embed_tokens.weight",
+    }
+
+    def __init__(self, config: FSMTConfig):
+        super().__init__(config)
+        self.encoder = FSMTEncoder(config)
+        self.decoder = FSMTDecoder(config)
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: torch.Tensor | None = None,
+        decoder_input_ids: torch.LongTensor | None = None,
+        decoder_attention_mask: torch.BoolTensor | None = None,
+        encoder_outputs: tuple[torch.FloatTensor] | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        decoder_inputs_embeds: torch.FloatTensor | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor] | Seq2SeqModelOutput:
+        r"""
+        decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+            Indices of decoder input sequence tokens in the vocabulary.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+            FSMT uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values`
+            is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`).
+        decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+            Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+            be used by default.
+        """
+        if decoder_input_ids is None:
+            use_cache = False
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        # make masks if user doesn't supply
+        if not use_cache and input_ids is not None:
+            decoder_input_ids, decoder_padding_mask, causal_mask = _prepare_fsmt_decoder_inputs(
+                self.config,
+                input_ids,
+                decoder_input_ids=decoder_input_ids,
+                decoder_padding_mask=decoder_attention_mask,
+                causal_mask_dtype=self.decoder.embed_tokens.weight.dtype,
+            )
+        else:
+            decoder_padding_mask, causal_mask = None, None
+
+        if decoder_input_ids is None and decoder_inputs_embeds is None:
+            raise ValueError("Make sure that `decoder_input_ids` or `decoder_inputs_embeds` are passed.")
+
+        if use_cache and past_key_values is None:
+            past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+
+        if encoder_outputs is None:
+            encoder_outputs = self.encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+        # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=False
+        elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
+            encoder_outputs = BaseModelOutput(
+                last_hidden_state=encoder_outputs[0],
+                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
+            )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        decoder_outputs = self.decoder(
+            decoder_input_ids,
+            encoder_outputs[0],
+            attention_mask,
+            decoder_padding_mask,
+            decoder_causal_mask=causal_mask,
+            inputs_embeds=decoder_inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        if not return_dict:
+            return decoder_outputs + encoder_outputs
+
+        return Seq2SeqModelOutput(
+            last_hidden_state=decoder_outputs.last_hidden_state,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+        )
+
+    def get_input_embeddings(self):
+        return self.encoder.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.encoder.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.decoder.embed_tokens
+
+    def set_output_embeddings(self, value):
+        self.decoder.embed_tokens = value
+
+
+@auto_docstring(
+    custom_intro="""
+    The FSMT Model with a language modeling head. Can be used for summarization.
+    """
+)
+class FSMTForConditionalGeneration(PretrainedFSMTModel, GenerationMixin):
+    base_model_prefix = "model"
+
+    def __init__(self, config: FSMTConfig):
+        super().__init__(config)
+        base_model = FSMTModel(config)
+        self.model = base_model
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        decoder_input_ids: torch.LongTensor | None = None,
+        decoder_attention_mask: torch.BoolTensor | None = None,
+        encoder_outputs: tuple[torch.FloatTensor] | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        decoder_inputs_embeds: torch.Tensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor] | Seq2SeqLMOutput:
+        r"""
+        decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+            Indices of decoder input sequence tokens in the vocabulary.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+            FSMT uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values`
+            is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`).
+        decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+            Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+            be used by default.
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example Translation:
+
+        ```python
+        >>> from transformers import AutoTokenizer, FSMTForConditionalGeneration
+
+        >>> mname = "facebook/wmt19-ru-en"
+        >>> model = FSMTForConditionalGeneration.from_pretrained(mname)
+        >>> tokenizer = AutoTokenizer.from_pretrained(mname)
+
+        >>> src_text = "Машинное обучение - это здорово, не так ли?"
+        >>> input_ids = tokenizer(src_text, return_tensors="pt").input_ids
+        >>> outputs = model.generate(input_ids, num_beams=5, num_return_sequences=3)
+        >>> tokenizer.decode(outputs[0], skip_special_tokens=True)
+        "Machine learning is great, isn't it?"
+        ```
+        """
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        if labels is not None:
+            use_cache = False
+
+        outputs = self.model(
+            input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            decoder_input_ids=decoder_input_ids,
+            decoder_inputs_embeds=decoder_inputs_embeds,
+            encoder_outputs=encoder_outputs,
+            decoder_attention_mask=decoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        lm_logits = outputs[0]
+
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # TODO(SS): do we need to ignore pad tokens in labels?
+            masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.tgt_vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (lm_logits,) + outputs[1:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+        return Seq2SeqLMOutput(
+            loss=masked_lm_loss,
+            logits=lm_logits,
+            past_key_values=outputs.past_key_values,
+            decoder_hidden_states=outputs.decoder_hidden_states,
+            decoder_attentions=outputs.decoder_attentions,
+            cross_attentions=outputs.cross_attentions,
+            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+            encoder_hidden_states=outputs.encoder_hidden_states,
+            encoder_attentions=outputs.encoder_attentions,
+        )
+
+    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
+        return shift_tokens_right(labels, self.config.pad_token_id)
+
+    def get_output_embeddings(self):
+        return self.model.decoder.embed_tokens
+
+    def set_output_embeddings(self, value):
+        self.model.decoder.embed_tokens = value
+
+
+class SinusoidalPositionalEmbedding(nn.Embedding):
+    """
+    This module produces sinusoidal positional embeddings of any length.
+
+    We don't want to save the weight of this embedding since it's not trained (deterministic) and it can be huge.
+
+    Padding symbols are ignored.
+
+    These embeddings get automatically extended in forward if more positions is needed.
+    """
+
+    def __init__(self, num_positions, embedding_dim, padding_idx):
+        super().__init__(num_positions, embedding_dim, padding_idx)
+
+    def make_weight(self, num_positions, embedding_dim, padding_idx):
+        weight = self.get_embedding(num_positions, embedding_dim, padding_idx)
+        # in forward put the weights on the correct dtype and device of the param
+        weight = weight.to(dtype=self.weight.dtype, device=self.weight.device)
+        self.weight = nn.Parameter(weight)
+        self.weight.detach_()
+        self.weight.requires_grad = False
+
+    @staticmethod
+    def get_embedding(num_embeddings, embedding_dim, padding_idx):
+        """
+        Build sinusoidal embeddings.
+
+        This matches the implementation in tensor2tensor, but differs slightly from the description in Section 3.5 of
+        "Attention Is All You Need".
+        """
+        half_dim = embedding_dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=torch.int64).float() * -emb)
+        emb = torch.arange(num_embeddings, dtype=torch.int64).float().unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
+        if embedding_dim % 2 == 1:
+            # zero pad
+            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
+        if padding_idx is not None:
+            emb[padding_idx, :] = 0
+        return emb
+
+    @staticmethod
+    def make_positions(tensor, padding_idx: int):
+        """
+        Replace non-padding symbols with their position numbers.
+
+        Position numbers begin at padding_idx+1. Padding symbols are ignored.
+        """
+        # The series of casts and type-conversions here are carefully
+        # balanced to both work with ONNX export and XLA. In particular XLA
+        # prefers ints, cumsum defaults to output longs, and ONNX doesn't know
+        # how to handle the dtype kwarg in cumsum.
+        mask = tensor.ne(padding_idx).int()
+        return (torch.cumsum(mask, dim=1).type_as(mask) * mask).long() + padding_idx
+
+    def forward(
+        self,
+        input,
+        incremental_state: Any | None = None,
+        timestep: Tensor | None = None,
+    ):
+        """Input is expected to be of size [bsz x seqlen]."""
+        bsz, seq_len = input.shape[:2]
+        max_pos = self.padding_idx + 1 + seq_len
+        if max_pos > self.weight.size(0):
+            # expand embeddings if needed
+            self.make_weight(max_pos, self.embedding_dim, self.padding_idx)
+        positions = self.make_positions(input, self.padding_idx)
+        return super().forward(positions)
+
+
+__all__ = ["FSMTForConditionalGeneration", "FSMTModel", "PretrainedFSMTModel"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_vl_moe/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_qwen3_vl_moe import *
+    from .modeling_qwen3_vl_moe import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py

@@ -0,0 +1,1886 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_qwen3_vl_moe.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import itertools
+import warnings
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import (
+    use_experts_implementation,
+    use_kernel_forward_from_hub,
+    use_kernel_func_from_hub,
+    use_kernelized_func,
+)
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput, MoeModelOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
+from ...utils.generic import (
+    accepts_precomputed_kwargs,
+    is_flash_attention_requested,
+    maybe_autocast,
+    merge_with_config_defaults,
+)
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from ...vision_utils import get_vision_bilinear_indices_and_weights, get_vision_cu_seqlens, get_vision_position_ids
+from .configuration_qwen3_vl_moe import Qwen3VLMoeConfig, Qwen3VLMoeTextConfig, Qwen3VLMoeVisionConfig
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Qwen3VLMoeTextRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        Qwen3VLMoeTextRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+@use_experts_implementation
+class Qwen3VLMoeTextExperts(nn.Module):
+    """Collection of expert weights stored as 3D tensors."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.hidden_dim = config.hidden_size
+        self.intermediate_dim = config.moe_intermediate_size
+        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
+        self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_index: torch.Tensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        final_hidden_states = torch.zeros_like(hidden_states)
+        with torch.no_grad():
+            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
+            expert_mask = expert_mask.permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+        for expert_idx in expert_hit:
+            expert_idx = expert_idx[0]
+            if expert_idx == self.num_experts:
+                continue
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            current_state = hidden_states[token_idx]
+            gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
+            current_hidden_states = self.act_fn(gate) * up
+            current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
+            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
+
+        return final_hidden_states
+
+
+class Qwen3VLMoeTextTopKRouter(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+        self.hidden_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
+
+    def forward(self, hidden_states):
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = F.linear(hidden_states, self.weight)  # (seq_len, num_experts)
+        router_probs = torch.nn.functional.softmax(router_logits, dtype=torch.float, dim=-1)
+        router_top_value, router_indices = torch.topk(router_probs, self.top_k, dim=-1)  # (seq_len, top_k)
+        router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
+        router_top_value = router_top_value.to(router_logits.dtype)
+        router_scores = router_top_value
+        return router_logits, router_scores, router_indices
+
+
+class Qwen3VLMoeTextSparseMoeBlock(nn.Module):
+    def __init__(self, config: Qwen3VLMoeTextConfig):
+        super().__init__()
+        self.experts = Qwen3VLMoeTextExperts(config)
+        self.gate = Qwen3VLMoeTextTopKRouter(config)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states_reshaped = hidden_states.view(-1, hidden_dim)
+        _, routing_weights, selected_experts = self.gate(hidden_states_reshaped)
+        final_hidden_states = self.experts(hidden_states_reshaped, selected_experts, routing_weights)
+        return final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+@use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class Qwen3VLMoeTextAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: Qwen3VLMoeTextConfig, layer_idx: int):
+        super().__init__()
+        self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+        self.q_norm = Qwen3VLMoeTextRMSNorm(
+            self.head_dim, eps=config.rms_norm_eps
+        )  # unlike olmo, only on the head dim!
+        self.k_norm = Qwen3VLMoeTextRMSNorm(
+            self.head_dim, eps=config.rms_norm_eps
+        )  # thus post q_norm does not need reshape
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class Qwen3VLMoeTextMLP(nn.Module):
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class Qwen3VLMoeTextDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: Qwen3VLMoeTextConfig, layer_idx: int):
+        super().__init__()
+        self.self_attn = Qwen3VLMoeTextAttention(config, layer_idx)
+        if (layer_idx not in config.mlp_only_layers) and (
+            config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
+        ):
+            self.mlp = Qwen3VLMoeTextSparseMoeBlock(config)
+        else:
+            self.mlp = Qwen3VLMoeTextMLP(config, intermediate_size=config.intermediate_size)
+        self.input_layernorm = Qwen3VLMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3VLMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.hidden_size = config.hidden_size
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring
+class Qwen3VLMoePreTrainedModel(PreTrainedModel):
+    config: Qwen3VLMoeConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen3VLMoeTextDecoderLayer", "Qwen3VLMoeVisionBlock"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, index=0),
+        "hidden_states": Qwen3VLMoeTextDecoderLayer,
+        "attentions": Qwen3VLMoeTextAttention,
+    }
+    input_modalities = ("text", "image", "video")
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        """Initialize the weights."""
+        super()._init_weights(module)
+        if hasattr(self.config, "initializer_range"):
+            std = self.config.initializer_range
+        else:
+            std = getattr(self.config.get_text_config(), "initializer_range", 0.02)
+        if isinstance(module, Qwen3VLMoeTextExperts):
+            init.normal_(module.gate_up_proj, mean=0.0, std=std)
+            init.normal_(module.down_proj, mean=0.0, std=std)
+        elif isinstance(module, Qwen3VLMoeTextTopKRouter):
+            init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, Qwen3VLMoeVisionRotaryEmbedding):
+            inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim))
+            init.copy_(module.inv_freq, inv_freq)
+
+
+class Qwen3VLMoeVisionRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, dim: int, theta: float = 10000.0) -> None:
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
+        return (position_ids.unsqueeze(-1) * self.inv_freq).flatten(1)
+
+
+def apply_rotary_pos_emb_vision(
+    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    orig_q_dtype = q.dtype
+    orig_k_dtype = k.dtype
+    q, k = q.float(), k.float()
+    cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    q_embed = q_embed.to(orig_q_dtype)
+    k_embed = k_embed.to(orig_k_dtype)
+    return q_embed, k_embed
+
+
+class Qwen3VLMoeVisionAttention(nn.Module):
+    def __init__(self, config: Qwen3VLMoeVisionConfig) -> None:
+        super().__init__()
+        self.dim = config.hidden_size
+        self.num_heads = config.num_heads
+        self.head_dim = self.dim // self.num_heads
+        self.num_key_value_groups = 1  # needed for eager attention
+        self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)
+        self.proj = nn.Linear(self.dim, self.dim)
+        self.scaling = self.head_dim**-0.5
+        self.config = config
+        self.attention_dropout = 0.0
+        self.is_causal = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rotary_pos_emb: torch.Tensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        seq_length = hidden_states.shape[0]
+        query_states, key_states, value_states = (
+            self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+        )
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
+
+        query_states = query_states.transpose(0, 1).unsqueeze(0)
+        key_states = key_states.transpose(0, 1).unsqueeze(0)
+        value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        if is_flash_attention_requested(self.config):
+            # Flash Attention: Use cu_seqlens for variable length attention
+            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+            attn_output, _ = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask=None,
+                scaling=self.scaling,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                cu_seq_lens_q=cu_seqlens,
+                cu_seq_lens_k=cu_seqlens,
+                max_length_q=max_seqlen,
+                max_length_k=max_seqlen,
+                is_causal=False,
+                **kwargs,
+            )
+        else:
+            # Other implementations: Process each chunk separately
+            lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+            splits = [
+                torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+            ]
+
+            attn_outputs = [
+                attention_interface(
+                    self,
+                    q,
+                    k,
+                    v,
+                    attention_mask=None,
+                    scaling=self.scaling,
+                    dropout=0.0 if not self.training else self.attention_dropout,
+                    is_causal=False,
+                    **kwargs,
+                )[0]
+                for q, k, v in zip(*splits)
+            ]
+            attn_output = torch.cat(attn_outputs, dim=1)
+
+        attn_output = attn_output.reshape(seq_length, -1).contiguous()
+        attn_output = self.proj(attn_output)
+        return attn_output
+
+
+class Qwen3VLMoeVisionMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
+        self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_state):
+        return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))
+
+
+class Qwen3VLMoeVisionBlock(GradientCheckpointingLayer):
+    def __init__(self, config, attn_implementation: str = "sdpa") -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+        self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+        self.attn = Qwen3VLMoeVisionAttention(config=config)
+        self.mlp = Qwen3VLMoeVisionMLP(config=config)
+
+    @auto_docstring
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rotary_pos_emb: torch.Tensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        r"""
+        cu_seqlens (`torch.Tensor`):
+            Cumulative sequence lengths used for packed variable-length attention in Flash Attention kernels.
+        rotary_pos_emb (`torch.Tensor`, *optional*):
+            Precomputed rotary positional embeddings applied to the vision attention query/key states.
+        """
+        hidden_states = hidden_states + self.attn(
+            self.norm1(hidden_states),
+            cu_seqlens=cu_seqlens,
+            rotary_pos_emb=rotary_pos_emb,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+        return hidden_states
+
+
+@auto_docstring
+@dataclass
+class BaseModelOutputWithDeepstackFeatures(BaseModelOutputWithPooling):
+    r"""
+    deepstack_features (`List[torch.FloatTensor]`, *optional*):
+        List of hidden-states (feature maps) from deepstack layers.
+    """
+
+    deepstack_features: list[torch.FloatTensor] | None = None
+
+
+class Qwen3VLMoeVisionPatchEmbed(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.patch_size = config.patch_size
+        self.temporal_patch_size = config.temporal_patch_size
+        self.in_channels = config.in_channels
+        self.embed_dim = config.hidden_size
+
+        kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
+        self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        target_dtype = self.proj.weight.dtype
+        hidden_states = hidden_states.view(
+            -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
+        )
+        hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+        return hidden_states
+
+
+class Qwen3VLMoeVisionPatchMerger(nn.Module):
+    def __init__(self, config: Qwen3VLMoeVisionConfig, use_postshuffle_norm=False) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)
+        self.use_postshuffle_norm = use_postshuffle_norm
+        self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6)
+        self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)
+        self.act_fn = nn.GELU()
+        self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size)
+        x = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
+        return x
+
+
+class Qwen3VLMoeVisionModel(Qwen3VLMoePreTrainedModel):
+    config: Qwen3VLMoeVisionConfig
+    input_modalities = ("image", "video")
+    _no_split_modules = ["Qwen3VLMoeVisionBlock"]
+    _can_record_outputs = {
+        "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, layer_name="mlp.gate", index=0),
+        "hidden_states": Qwen3VLMoeVisionBlock,
+        "attentions": Qwen3VLMoeVisionAttention,
+    }
+
+    def __init__(self, config, *inputs, **kwargs) -> None:
+        super().__init__(config, *inputs, **kwargs)
+        self.spatial_merge_size = config.spatial_merge_size
+        self.patch_size = config.patch_size
+        self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
+
+        self.patch_embed = Qwen3VLMoeVisionPatchEmbed(
+            config=config,
+        )
+
+        self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size)
+        self.num_grid_per_side = int(config.num_position_embeddings**0.5)
+
+        head_dim = config.hidden_size // config.num_heads
+        self.rotary_pos_emb = Qwen3VLMoeVisionRotaryEmbedding(head_dim // 2)
+
+        self.blocks = nn.ModuleList([Qwen3VLMoeVisionBlock(config) for _ in range(config.depth)])
+        self.merger = Qwen3VLMoeVisionPatchMerger(
+            config=config,
+            use_postshuffle_norm=False,
+        )
+
+        self.deepstack_visual_indexes = config.deepstack_visual_indexes
+        self.deepstack_merger_list = nn.ModuleList(
+            [
+                Qwen3VLMoeVisionPatchMerger(
+                    config=config,
+                    use_postshuffle_norm=True,
+                )
+                for _ in range(len(config.deepstack_visual_indexes))
+            ]
+        )
+
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
+        warnings.warn(
+            f"`{self.__class__.__name__}.rot_pos_emb` is deprecated and will be removed in v5.11. Use `get_vision_position_ids` from `transformers.vision_utils` and apply the rotary embedding module.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size)
+        rotary_pos_emb = self.rotary_pos_emb(position_ids)
+        return rotary_pos_emb
+
+    def fast_pos_embed_interpolate(self, grid_thw):
+        warnings.warn(
+            f"`{self.__class__.__name__}.fast_pos_embed_interpolate` is deprecated and will be removed in v5.11. Use `get_vision_bilinear_indices_and_weights` from `transformers.vision_utils` and apply `self.pos_embed`.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        bilinear_indices, bilinear_weights = get_vision_bilinear_indices_and_weights(
+            grid_thw,
+            num_grid_per_side=self.num_grid_per_side,
+            spatial_merge_size=self.config.spatial_merge_size,
+        )
+        return (self.pos_embed(bilinear_indices) * bilinear_weights[:, :, None]).sum(0)
+
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+    ) -> tuple | BaseModelOutputWithDeepstackFeatures:
+        """
+        Args:
+            hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
+                The final hidden states of the model.
+            grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+                The temporal, height and width of feature shape of each image in LLM.
+
+        Returns:
+            `torch.Tensor`: hidden_states.
+        """
+        bilinear_indices, bilinear_weights = get_vision_bilinear_indices_and_weights(
+            grid_thw,
+            num_grid_per_side=self.num_grid_per_side,
+            spatial_merge_size=self.config.spatial_merge_size,
+            kwargs=kwargs,
+        )
+        position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size, kwargs=kwargs)
+        cu_seqlens = get_vision_cu_seqlens(grid_thw, kwargs=kwargs)
+
+        hidden_states = self.patch_embed(hidden_states)
+        pos_embeds = (self.pos_embed(bilinear_indices) * bilinear_weights[:, :, None]).sum(0)
+        hidden_states = hidden_states + pos_embeds.to(hidden_states.dtype)
+        rotary_pos_emb = self.rotary_pos_emb(position_ids)
+
+        seq_len, _ = hidden_states.size()
+        hidden_states = hidden_states.reshape(seq_len, -1)
+        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+        emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+        position_embeddings = (emb.cos(), emb.sin())
+
+        deepstack_feature_lists = []
+        for layer_num, blk in enumerate(self.blocks):
+            hidden_states = blk(
+                hidden_states,
+                cu_seqlens=cu_seqlens,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+            if layer_num in self.deepstack_visual_indexes:
+                deepstack_feature = self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)](
+                    hidden_states
+                )
+                deepstack_feature_lists.append(deepstack_feature)
+
+        merged_hidden_states = self.merger(hidden_states)
+
+        return BaseModelOutputWithDeepstackFeatures(
+            last_hidden_state=hidden_states,
+            pooler_output=merged_hidden_states,
+            deepstack_features=deepstack_feature_lists,
+        )
+
+
+class Qwen3VLMoeTextRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: Qwen3VLMoeTextConfig, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+        self.mrope_section = config.rope_parameters.get("mrope_section", [24, 20, 20])
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: Qwen3VLMoeTextConfig | None = None,
+        device: Optional["torch.device"] = None,
+        seq_len: int | None = None,
+    ) -> tuple["torch.Tensor", float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters["rope_theta"]
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        # Compute the inverse frequencies
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, attention_factor
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        # In contrast to other models, Qwen3VLMoe has different position ids for the grids
+        # So we expand the inv_freq to shape (3, ...)
+        if position_ids.ndim == 2:
+            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+        inv_freq_expanded = (
+            self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1).to(x.device)
+        )
+        position_ids_expanded = position_ids[:, :, None, :].float()  # shape (3, bs, 1, positions)
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+            freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+    def apply_interleaved_mrope(self, freqs, mrope_section):
+        """Apply interleaved MRoPE to 3D rotary embeddings.
+        Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
+        interleaved [THWTHWTHW...TT], preserving frequency continuity.
+        args:
+            x: (3, bs, seq_len, head_dim // 2)
+            mrope_section: (3,)
+        returns:
+            x_t: (bs, seq_len, head_dim // 2)
+        """
+        freqs_t = freqs[0]  # just overwrite the first dimension T
+        for dim, offset in enumerate((1, 2), start=1):  # H, W
+            length = mrope_section[dim] * 3
+            idx = slice(offset, length, 3)
+            freqs_t[..., idx] = freqs[dim, ..., idx]
+        return freqs_t
+
+
+@auto_docstring(
+    custom_intro=(
+        "Text part of Qwen3VLMoe, "
+        "not a pure text-only model, as DeepStack integrates visual features into the early hidden states."
+    )
+)
+class Qwen3VLMoeTextModel(Qwen3VLMoePreTrainedModel):
+    config: Qwen3VLMoeTextConfig
+    input_modalities = ("text",)
+    _no_split_modules = ["Qwen3VLMoeTextDecoderLayer"]
+
+    def __init__(self, config: Qwen3VLMoeTextConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [Qwen3VLMoeTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = Qwen3VLMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen3VLMoeTextRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        # args for deepstack
+        visual_pos_masks: torch.Tensor | None = None,
+        deepstack_visual_embeds: list[torch.Tensor] | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple | MoeModelOutputWithPast:
+        r"""
+        visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*):
+            The mask of the visual positions.
+        deepstack_visual_embeds (`list[torch.Tensor]`, *optional*):
+            The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim).
+            The feature is extracted from the different visual encoder layers, and fed to the decoder
+            hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334).
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        # torch.jit.trace() doesn't support cache objects in the output
+        if use_cache and past_key_values is None and not torch.jit.is_tracing():
+            past_key_values = DynamicCache(config=self.config)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # the hard coded `4` is for text, temporal, height and width.
+        if position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1)
+        elif position_ids.ndim == 2:
+            position_ids = position_ids[None, ...].expand(4, position_ids.shape[0], -1)
+
+        if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+            text_position_ids = position_ids[0]
+            position_ids = position_ids[1:]
+        else:
+            text_position_ids = None
+
+        attention_mask = create_causal_mask(
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_ids=text_position_ids,
+        )
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        for layer_idx, decoder_layer in enumerate(self.layers):
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=text_position_ids,
+                past_key_values=past_key_values,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+            hidden_states = layer_outputs
+
+            # add visual features to the hidden states of first several layers
+            if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
+                hidden_states = self._deepstack_process(
+                    hidden_states,
+                    visual_pos_masks,
+                    deepstack_visual_embeds[layer_idx],
+                )
+
+        hidden_states = self.norm(hidden_states)
+
+        return MoeModelOutputWithPast(  # only diff with Qwen3VLTextModel
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+    def _deepstack_process(
+        self, hidden_states: torch.Tensor, visual_pos_masks: torch.Tensor, visual_embeds: torch.Tensor
+    ):
+        visual_pos_masks = visual_pos_masks.to(hidden_states.device)
+        visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)
+        hidden_states = hidden_states.clone()
+        local_this = hidden_states[visual_pos_masks, :] + visual_embeds
+        hidden_states[visual_pos_masks, :] = local_this
+        return hidden_states
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for Llava outputs, with hidden states and attentions.
+    """
+)
+@dataclass
+class Qwen3VLMoeModelOutputWithPast(ModelOutput):
+    r"""
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+        The rope index difference between sequence length and multimodal rope.
+    """
+
+    last_hidden_state: torch.FloatTensor | None = None
+    past_key_values: Cache | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+    rope_deltas: torch.LongTensor | None = None
+    router_logits: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for Qwen3VLMoe causal language model (or autoregressive) outputs.
+    """
+)
+@dataclass
+class Qwen3VLMoeCausalLMOutputWithPast(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Language modeling loss (for next-token prediction).
+    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+        The rope index difference between sequence length and multimodal rope.
+    """
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    past_key_values: Cache | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+    rope_deltas: torch.LongTensor | None = None
+    router_logits: tuple[torch.FloatTensor] | None = None
+    aux_loss: torch.FloatTensor | None = None
+
+
+@auto_docstring
+class Qwen3VLMoeModel(Qwen3VLMoePreTrainedModel):
+    base_model_prefix = "model"
+    # Reference: fix gemma3 grad acc #37208
+    accepts_loss_kwargs = False
+    config: Qwen3VLMoeConfig
+    _no_split_modules = ["Qwen3VLMoeTextDecoderLayer", "Qwen3VLMoeVisionBlock"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.visual = Qwen3VLMoeVisionModel._from_config(config.vision_config)
+        self.language_model = Qwen3VLMoeTextModel._from_config(config.text_config)
+        self.rope_deltas = None  # cache rope_deltas here
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_vision_position_ids(
+        self,
+        start_position: int,
+        grid_thw: list[int, int, int] | torch.Tensor,
+        temp_merge_size: int = 1,
+        spatial_merge_size: int = 1,
+        time_interval: int = 1,
+        device: str | torch.device | None = None,
+    ):
+        """
+        Compute 3D positional indices for vision tokens derived from a single image or video input.
+
+        The positions are generated from the input grid defined by temporal (T), height (H), and
+        width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the
+        merge sizes used in the vision backbone. The resulting positions are offset by `start_position`.
+
+        Args:
+            start_position (`int`):
+                Offset added to all computed positional indices.
+            grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`):
+                The (T, H, W) grid representing the feature layout of the current image or video after patch embedding.
+            temp_merge_size (`int`, *optional*):
+                Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided
+                by this value. Defaults to 1.
+            spatial_merge_size (`int`, *optional*):
+                Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided
+                by this value. Defaults to 1.
+            time_interval (`int`, *optional*):
+                Spacing factor applied between consecutive temporal position indices.Defaults to 1.
+            device (`str` or `torch.device`, *optional*):
+                Device on which the resulting tensor is allocated. If `None`, uses the current default device.
+
+        Returns:
+            torch.LongTensor of shape (3, sequence_length):
+                Positional indices for temporal, height, and width dimensions,
+                flattened into sequence form and offset by `start_position`.
+        """
+        llm_grid_t, llm_grid_h, llm_grid_w = (
+            grid_thw[0].item() // temp_merge_size,
+            grid_thw[1].item() // spatial_merge_size,
+            grid_thw[2].item() // spatial_merge_size,
+        )
+
+        # Add `start_position` after arange for compile
+        position_temporal = torch.arange(llm_grid_t, device=device) * time_interval
+        position_width = torch.arange(llm_grid_w, device=device) + start_position
+        position_height = torch.arange(llm_grid_h, device=device) + start_position
+
+        # Repeat the positions per each grid and per video frame. Repeat patterns are important
+        # do not modify without checking values!
+        position_width = position_width.repeat(llm_grid_h * llm_grid_t)
+        position_height = position_height.repeat_interleave(llm_grid_w).repeat(llm_grid_t)
+        # Important: add `start_positions` after applying `time_interval`, order matters
+        position_temporal = position_temporal.repeat_interleave(llm_grid_h * llm_grid_w) + start_position
+        vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0)
+
+        return vision_position_ids
+
+    def get_rope_index(
+        self,
+        input_ids: torch.LongTensor,
+        mm_token_type_ids: torch.IntTensor,
+        image_grid_thw: torch.LongTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Difference from Qwen2VL/Qwen2.5VL's get_rope_index:
+        - Since Qwen3.5 use timestamps to separate videos, like <t1> <vision_start> <frame1> <vision_end> <t2> <vision_start> <frame2> <vision_end>, the video_grid_thw should also be split too.
+
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+                it.
+            mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`):
+                Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2).
+            image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+                The temporal, height and width of feature shape of each image in LLM.
+            video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+                The temporal, height and width of feature shape of each video in LLM.
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+        Returns:
+            position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
+            mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
+        """
+
+        # Separate video grid thw into multiple grids because timestamps are used to separate videos.
+        if video_grid_thw is not None:
+            video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
+            video_grid_thw[:, 0] = 1
+        spatial_merge_size = self.config.vision_config.spatial_merge_size
+
+        mrope_position_deltas = []
+        position_ids = torch.zeros(
+            3,
+            input_ids.shape[0],
+            input_ids.shape[1],
+            dtype=input_ids.dtype,
+            device=input_ids.device,
+        )
+        grid_iters = {
+            1: iter(image_grid_thw) if image_grid_thw is not None else None,
+            2: iter(video_grid_thw) if video_grid_thw is not None else None,
+        }
+
+        for batch_idx, current_input_ids in enumerate(input_ids):
+            input_token_type = mm_token_type_ids[batch_idx]
+            if attention_mask is not None:
+                current_input_ids = current_input_ids[attention_mask[batch_idx].bool()]
+                input_token_type = input_token_type[attention_mask[batch_idx].bool()]
+
+            input_type_group = []
+            for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]):
+                group = list(group)
+                start_index = group[0][0]
+                end_index = group[-1][0] + 1
+                input_type_group.append((key, start_index, end_index))
+
+            current_pos = 0
+            llm_pos_ids_list = []
+            for modality_type, start_idx, end_idx in input_type_group:
+                # text == 0
+                if modality_type == 0:
+                    text_len = end_idx - start_idx
+                    llm_pos_ids_list.append(
+                        torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos
+                    )
+                    current_pos += text_len
+                # image == 1, video == 2
+                else:
+                    grid_thw = next(grid_iters[modality_type])
+                    vision_position_ids = self.get_vision_position_ids(
+                        current_pos, grid_thw, 1, spatial_merge_size, device=input_ids.device
+                    )
+                    llm_pos_ids_list.append(vision_position_ids)
+                    current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size
+            llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+            if attention_mask is not None:
+                position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device)
+            else:
+                position_ids[:, batch_idx] = llm_positions.to(position_ids.device)
+            mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids))
+        mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+        return position_ids, mrope_position_deltas
+
+    @accepts_precomputed_kwargs(modality="video")
+    @can_return_tuple
+    @auto_docstring
+    def get_video_features(
+        self,
+        pixel_values_videos: torch.FloatTensor,
+        video_grid_thw: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithDeepstackFeatures:
+        r"""
+        pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+            The tensors corresponding to the input videos.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            The temporal, height and width of feature shape of each video in LLM.
+        """
+        # Same implementation as for images
+        return self.get_image_features(pixel_values_videos, video_grid_thw, **kwargs)
+
+    @accepts_precomputed_kwargs(modality="image")
+    @can_return_tuple
+    @auto_docstring
+    def get_image_features(
+        self,
+        pixel_values: torch.FloatTensor,
+        image_grid_thw: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithDeepstackFeatures:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+            The tensors corresponding to the input images.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            The temporal, height and width of feature shape of each image in LLM.
+        """
+        pixel_values = pixel_values.type(self.visual.dtype)
+        vision_output: BaseModelOutputWithDeepstackFeatures = self.visual(
+            pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs
+        )
+        image_embeds = vision_output.pooler_output
+        split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
+        image_embeds = torch.split(image_embeds, split_sizes)
+        vision_output.pooler_output = image_embeds
+
+        return vision_output
+
+    def get_placeholder_mask(
+        self,
+        input_ids: torch.LongTensor,
+        inputs_embeds: torch.FloatTensor,
+        image_features: torch.FloatTensor | None = None,
+        video_features: torch.FloatTensor | None = None,
+    ):
+        """
+        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+        equal to the length of multimodal features. If the lengths are different, an error is raised.
+        """
+        if input_ids is None:
+            special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_image_mask = special_image_mask.all(-1)
+            special_video_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_video_mask = special_video_mask.all(-1)
+        else:
+            special_image_mask = input_ids == self.config.image_token_id
+            special_video_mask = input_ids == self.config.video_token_id
+
+        n_image_tokens = special_image_mask.sum()
+        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+        if image_features is not None:
+            torch_compilable_check(
+                inputs_embeds[special_image_mask].numel() == image_features.numel(),
+                f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}",
+            )
+
+        n_video_tokens = special_video_mask.sum()
+        special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+        if video_features is not None:
+            torch_compilable_check(
+                inputs_embeds[special_video_mask].numel() == video_features.numel(),
+                f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}",
+            )
+        return special_image_mask, special_video_mask
+
+    def compute_3d_position_ids(
+        self,
+        input_ids: torch.Tensor | None,
+        inputs_embeds: torch.Tensor | None,
+        image_grid_thw: torch.Tensor | None = None,
+        video_grid_thw: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: torch.Tensor | None = None,
+        mm_token_type_ids: torch.IntTensor | None = None,
+    ) -> torch.Tensor | None:
+        past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+        has_multimodal = image_grid_thw is not None or video_grid_thw is not None
+        if has_multimodal and mm_token_type_ids is None and input_ids is not None:
+            raise ValueError(
+                "Multimodal data was passed (via `image_grid_thw` or `video_grid_thw`) but `mm_token_type_ids` is "
+                "missing. Please pass `mm_token_type_ids` to the model so that multimodal RoPE (M-RoPE) can be "
+                "computed correctly. `mm_token_type_ids` is returned by the processor alongside `input_ids`."
+            )
+        can_compute_mrope = input_ids is not None and mm_token_type_ids is not None and has_multimodal
+
+        if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0):
+            position_ids, rope_deltas = self.get_rope_index(
+                input_ids,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                attention_mask=attention_mask,
+                mm_token_type_ids=mm_token_type_ids,
+            )
+            self.rope_deltas = rope_deltas
+        # Use pre-calculated rope-deltas to infer correct 3D position ids during incremental
+        # generation (past_key_values_length > 0) or when only inputs_embeds is provided (no input_ids
+        # to recompute from). Skip when input_ids is provided without past_key_values to avoid shape
+        # mismatches from stale rope_deltas (e.g., training forward pass after generation).
+        elif self.rope_deltas is not None and (past_key_values_length > 0 or input_ids is None):
+            batch_size, seq_length, _ = inputs_embeds.shape
+            if attention_mask is not None:
+                position_ids = attention_mask.long().cumsum(-1) - 1
+                position_ids = position_ids.masked_fill(attention_mask == 0, 0)
+                position_ids = position_ids.view(1, batch_size, -1).repeat(3, 1, 1).to(inputs_embeds.device)
+            else:
+                position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length)
+                position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device)
+            delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0)
+            position_ids = position_ids + delta.to(device=inputs_embeds.device)
+        else:
+            # Can't build correct 3D positions. Let the model infer it
+            position_ids = None
+        return position_ids
+
+    @auto_docstring
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        pixel_values: torch.Tensor | None = None,
+        pixel_values_videos: torch.FloatTensor | None = None,
+        image_grid_thw: torch.LongTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        mm_token_type_ids: torch.IntTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | Qwen3VLMoeModelOutputWithPast:
+        r"""
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            The temporal, height and width of feature shape of each image in LLM.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            The temporal, height and width of feature shape of each video in LLM.
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        image_mask = None
+        video_mask = None
+
+        if pixel_values is not None:
+            image_outputs: BaseModelOutputWithDeepstackFeatures = self.get_image_features(
+                pixel_values, image_grid_thw, return_dict=True, **kwargs
+            )
+            image_embeds = image_outputs.pooler_output
+            deepstack_image_embeds = image_outputs.deepstack_features
+            image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+            image_mask, _ = self.get_placeholder_mask(
+                input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+            )
+            inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+        if pixel_values_videos is not None:
+            video_outputs: BaseModelOutputWithDeepstackFeatures = self.get_video_features(
+                pixel_values_videos, video_grid_thw, return_dict=True, **kwargs
+            )
+            video_embeds = video_outputs.pooler_output
+            deepstack_video_embeds = video_outputs.deepstack_features
+            video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+            _, video_mask = self.get_placeholder_mask(
+                input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
+            )
+            inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+        visual_pos_masks = None
+        deepstack_visual_embeds = None
+        if image_mask is not None and video_mask is not None:
+            # aggregate visual_pos_masks and deepstack_visual_embeds
+            image_mask = image_mask[..., 0]
+            video_mask = video_mask[..., 0]
+            visual_pos_masks = image_mask | video_mask
+            deepstack_visual_embeds = []
+            image_mask_joint = image_mask[visual_pos_masks]
+            video_mask_joint = video_mask[visual_pos_masks]
+            for img_embed, vid_embed in zip(deepstack_image_embeds, deepstack_video_embeds):
+                embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1]).to(img_embed.device)
+                embed_joint[image_mask_joint, :] = img_embed
+                embed_joint[video_mask_joint, :] = vid_embed
+                deepstack_visual_embeds.append(embed_joint)
+        elif image_mask is not None:
+            image_mask = image_mask[..., 0]
+            visual_pos_masks = image_mask
+            deepstack_visual_embeds = deepstack_image_embeds
+        elif video_mask is not None:
+            video_mask = video_mask[..., 0]
+            visual_pos_masks = video_mask
+            deepstack_visual_embeds = deepstack_video_embeds
+
+        if position_ids is None:
+            position_ids = self.compute_3d_position_ids(
+                input_ids=input_ids,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                mm_token_type_ids=mm_token_type_ids,
+            )
+
+        outputs = self.language_model(
+            input_ids=None,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            visual_pos_masks=visual_pos_masks,
+            deepstack_visual_embeds=deepstack_visual_embeds,
+            **kwargs,
+        )
+
+        return Qwen3VLMoeModelOutputWithPast(
+            **outputs,
+            rope_deltas=self.rope_deltas,
+        )
+
+
+def load_balancing_loss_func(
+    gate_logits: torch.Tensor | tuple[torch.Tensor] | None,
+    num_experts: int | None = None,
+    top_k=2,
+    attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor | int:
+    r"""
+    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+
+    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
+    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+    experts is too unbalanced.
+
+    Args:
+        gate_logits:
+            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
+            shape [batch_size X sequence_length, num_experts].
+        num_experts:
+            Number of experts
+        top_k:
+            The number of experts to route per-token, can be also interpreted as the `top-k` routing
+            parameter.
+        attention_mask (`torch.Tensor`, *optional*):
+            The attention_mask used in forward function
+            shape [batch_size X sequence_length] if not None.
+
+    Returns:
+        The auxiliary loss.
+    """
+    if gate_logits is None or not isinstance(gate_logits, tuple):
+        return 0
+
+    if isinstance(gate_logits, tuple):
+        compute_device = gate_logits[0].device
+        concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
+
+    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
+
+    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
+
+    if attention_mask is None:
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.mean(routing_weights, dim=0)
+    else:
+        batch_size, sequence_length = attention_mask.shape
+        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
+        expert_attention_mask = (
+            attention_mask[None, :, :, None, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
+            .reshape(-1, top_k, num_experts)
+            .to(compute_device)
+        )
+
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
+            expert_attention_mask, dim=0
+        )
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
+        router_per_expert_attention_mask = (
+            attention_mask[None, :, :, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
+            .reshape(-1, num_experts)
+            .to(compute_device)
+        )
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
+            router_per_expert_attention_mask, dim=0
+        )
+
+    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+    return overall_loss * num_experts
+
+
+class Qwen3VLMoeForConditionalGeneration(Qwen3VLMoePreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+    # Reference: fix gemma3 grad acc #37208
+    accepts_loss_kwargs = False
+    config: Qwen3VLMoeConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Qwen3VLMoeModel(config)
+        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+
+        self.post_init()
+
+    @auto_docstring
+    def get_video_features(
+        self,
+        pixel_values_videos: torch.FloatTensor,
+        video_grid_thw: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithDeepstackFeatures:
+        r"""
+        pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+            The tensors corresponding to the input videos.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            The temporal, height and width of feature shape of each video in LLM.
+        """
+        return self.model.get_video_features(pixel_values_videos, video_grid_thw, **kwargs)
+
+    @auto_docstring
+    def get_image_features(
+        self,
+        pixel_values: torch.FloatTensor,
+        image_grid_thw: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithDeepstackFeatures:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+            The tensors corresponding to the input images.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            The temporal, height and width of feature shape of each image in LLM.
+        """
+        return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs)
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        pixel_values: torch.Tensor | None = None,
+        pixel_values_videos: torch.FloatTensor | None = None,
+        image_grid_thw: torch.LongTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        mm_token_type_ids: torch.IntTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | Qwen3VLMoeCausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            The temporal, height and width of feature shape of each image in LLM.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            The temporal, height and width of feature shape of each video in LLM.
+
+        Example:
+        ```python
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+        >>> from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
+
+        >>> model = Qwen3VLMoeForConditionalGeneration.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto")
+        >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct")
+
+        >>> messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
+                    },
+                    {"type": "text", "text": "Describe this image in short."},
+                ],
+            }
+        ]
+
+        >>> # Preparation for inference
+        >>> inputs = processor.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt"
+        )
+        >>> inputs = inputs.to(model.device)
+
+        >>> # Generate
+        >>> generated_ids = model.generate(**inputs, max_new_tokens=128)
+        >>> generated_ids_trimmed = [
+            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        >>> processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "A woman in a plaid shirt sits on a sandy beach at sunset, smiling as she gives a high-five to a yellow Labrador Retriever wearing a harness. The ocean waves roll in the background."
+        ```"""
+
+        outputs = self.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            mm_token_type_ids=mm_token_type_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+
+        aux_loss = None
+        if kwargs.get("output_router_logits", False):
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits,
+                self.config.text_config.num_experts,
+                self.config.text_config.num_experts_per_tok,
+                attention_mask,
+            )
+            if labels is not None:
+                loss += self.config.text_config.router_aux_loss_coef * aux_loss.to(
+                    loss.device
+                )  # make sure to reside in the same device
+
+        return Qwen3VLMoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            rope_deltas=outputs.rope_deltas,
+            router_logits=outputs.router_logits,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        position_ids=None,
+        use_cache=True,
+        pixel_values=None,
+        pixel_values_videos=None,
+        image_grid_thw=None,
+        video_grid_thw=None,
+        is_first_iteration=False,
+        **kwargs,
+    ):
+        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            position_ids=position_ids,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            use_cache=use_cache,
+            is_first_iteration=is_first_iteration,
+            **kwargs,
+        )
+
+        if not is_first_iteration and use_cache:
+            model_inputs["pixel_values"] = None
+            model_inputs["pixel_values_videos"] = None
+
+        return model_inputs
+
+    def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs):
+        # Overwritten -- requires 3D position ids
+
+        text_positions = super()._prepare_position_ids_for_generation(inputs_tensor, model_kwargs)
+
+        # Early exit in case we are continuing generation from past kv
+        past_length = 0
+        if (cache := model_kwargs.get("past_key_values")) is not None:
+            past_length = cache.get_seq_length()
+        if past_length != 0 and self.model.rope_deltas is not None:
+            position_ids = text_positions[None, ...] + self.model.rope_deltas
+            return position_ids
+
+        # Otherwise compute 3d position ids for vision tokens and concat with text position ids
+        if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0:
+            inputs_tensor = model_kwargs["input_ids"]
+
+        is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long]
+        if (
+            is_input_ids
+            and model_kwargs.get("mm_token_type_ids") is not None
+            and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None)
+        ):
+            model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"}
+            vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs)
+            self.model.rope_deltas = rope_deltas
+        else:
+            vision_positions = text_positions.unsqueeze(0).expand(3, -1, -1)
+            self.model.rope_deltas = torch.zeros(
+                inputs_tensor.shape[0], 1, dtype=torch.long, device=inputs_tensor.device
+            )
+
+        # Concatenate "text + vision" positions into [4, bs, seq-len]
+        text_positions = text_positions[None, ...]
+        position_ids = torch.cat([text_positions, vision_positions], dim=0)
+
+        return position_ids
+
+    def _get_image_nums_and_video_nums(
+        self,
+        input_ids: torch.LongTensor | None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Get the number of images and videos for each sample to calculate the separation length of the sample tensor.
+        These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications.
+
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                Indices of input sequence tokens in the vocabulary.
+
+        Returns:
+            image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`)
+            video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`)
+        """
+        image_token_id = self.config.image_token_id
+        video_token_id = self.config.video_token_id
+        vision_start_token_id = self.config.vision_start_token_id
+
+        if inputs_embeds is not None:
+            vision_start_mask = (
+                inputs_embeds
+                == self.get_input_embeddings()(
+                    torch.tensor(vision_start_token_id, dtype=torch.long, device=inputs_embeds.device)
+                )
+            )[..., 0]
+            image_mask = (
+                inputs_embeds
+                == self.get_input_embeddings()(
+                    torch.tensor(image_token_id, dtype=torch.long, device=inputs_embeds.device)
+                )
+            )[..., 0]
+            video_mask = (
+                inputs_embeds
+                == self.get_input_embeddings()(
+                    torch.tensor(video_token_id, dtype=torch.long, device=inputs_embeds.device)
+                )
+            )[..., 0]
+        else:
+            vision_start_mask = input_ids == vision_start_token_id
+            image_mask = input_ids == image_token_id
+            video_mask = input_ids == video_token_id
+
+        vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1)
+        image_nums = torch.sum(vision_first_mask & image_mask, dim=1)
+        video_nums = torch.sum(vision_first_mask & video_mask, dim=1)
+
+        return image_nums, video_nums
+
+    def _expand_inputs_for_generation(
+        self,
+        expand_size: int = 1,
+        is_encoder_decoder: bool = False,
+        input_ids: torch.LongTensor | None = None,
+        **model_kwargs,
+    ) -> tuple[torch.LongTensor, dict[str, Any]]:
+        # Overwritten -- Qwen3VLMoe use timestamps and remove second_per_grid_ts
+        # Support for expanding tensors without a batch size dimension
+        # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw
+        # pixel_values.shape[0] is sum(seqlen_images for samples)
+        # image_grid_thw.shape[0] is sum(num_images for samples)
+
+        if expand_size == 1:
+            return input_ids, model_kwargs
+
+        visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw"]
+
+        def _expand_dict_for_generation_visual(dict_to_expand):
+            image_grid_thw = model_kwargs.get("image_grid_thw", None)
+            video_grid_thw = model_kwargs.get("video_grid_thw", None)
+            image_nums, video_nums = self._get_image_nums_and_video_nums(
+                input_ids, inputs_embeds=model_kwargs.get("inputs_embeds", None)
+            )
+
+            # video_nums: (batch_size,)
+            # since video_nums is the number of videos in the input dependent on the input_ids(vision_start),
+            # but Qwen3VLMoe append vision_start to each frame of each video, so we need to recover the real video_nums according to video_grid_thw
+            if video_grid_thw is not None:
+                cumulative_frame_counts = torch.cumsum(video_grid_thw[:, 0], dim=0)
+                cumulative_token_video_counts = torch.cumsum(video_nums, dim=0)
+                # Find video boundaries in cumulative_frame_counts
+                video_boundary_indices = torch.searchsorted(cumulative_frame_counts, cumulative_token_video_counts)
+                # example: video_boundary_indices = [3, 5] means video_nums = [4, 2]
+                video_nums = torch.diff(torch.cat([-video_boundary_indices.new_ones(1), video_boundary_indices]))
+
+            def _repeat_interleave_samples(x, lengths, repeat_times):
+                samples = torch.split(x, lengths)
+                repeat_args = [repeat_times] + [1] * (x.dim() - 1)
+                result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0)
+                return result
+
+            for key in dict_to_expand:
+                if key == "pixel_values":
+                    # split images into samples
+                    samples = torch.split(image_grid_thw, list(image_nums))
+                    # compute the sequence length of images for each sample
+                    lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+                    dict_to_expand[key] = _repeat_interleave_samples(
+                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+                    )
+                elif key == "image_grid_thw":
+                    # get the num of images for each sample
+                    lengths = list(image_nums)
+                    dict_to_expand[key] = _repeat_interleave_samples(
+                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+                    )
+                elif key == "pixel_values_videos":
+                    samples = torch.split(video_grid_thw, list(video_nums))
+                    lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+                    dict_to_expand[key] = _repeat_interleave_samples(
+                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+                    )
+                elif key == "video_grid_thw":
+                    lengths = list(video_nums)
+                    dict_to_expand[key] = _repeat_interleave_samples(
+                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+                    )
+            return dict_to_expand
+
+        def _expand_dict_for_generation(dict_to_expand):
+            for key in dict_to_expand:
+                if key == "position_ids" and dict_to_expand[key].ndim == 3:
+                    dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=1)
+                elif (
+                    dict_to_expand[key] is not None
+                    and isinstance(dict_to_expand[key], torch.Tensor)
+                    and key not in visual_keys
+                ):
+                    dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+            return dict_to_expand
+
+        model_kwargs = _expand_dict_for_generation_visual(model_kwargs)
+
+        if input_ids is not None:
+            input_ids = input_ids.repeat_interleave(expand_size, dim=0)
+
+        model_kwargs = _expand_dict_for_generation(model_kwargs)
+
+        if is_encoder_decoder:
+            if model_kwargs.get("encoder_outputs") is None:
+                raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
+            model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])
+
+        return input_ids, model_kwargs
+
+
+__all__ = [
+    "Qwen3VLMoeVisionModel",
+    "Qwen3VLMoeForConditionalGeneration",
+    "Qwen3VLMoeModel",
+    "Qwen3VLMoePreTrainedModel",
+    "Qwen3VLMoeTextModel",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py

@@ -0,0 +1,467 @@
+# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Qwen3-VL-MOE model."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub.dataclasses import strict
+
+from ... import initialization as init
+from ...cache_utils import Cache, DynamicCache
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_outputs import MoeModelOutputWithPast
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils.output_capturing import OutputRecorder
+from ..qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig
+from ..qwen3_moe.modeling_qwen3_moe import (
+    Qwen3MoeDecoderLayer,
+    Qwen3MoeExperts,
+    Qwen3MoePreTrainedModel,
+    Qwen3MoeRMSNorm,
+    Qwen3MoeSparseMoeBlock,
+    load_balancing_loss_func,
+)
+from ..qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig
+from ..qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLCausalLMOutputWithPast,
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLModelOutputWithPast,
+    Qwen3VLTextAttention,
+    Qwen3VLTextModel,
+    Qwen3VLVisionAttention,
+    Qwen3VLVisionBlock,
+    Qwen3VLVisionModel,
+    Qwen3VLVisionRotaryEmbedding,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct")
+@strict
+class Qwen3VLMoeTextConfig(Qwen3MoeConfig):
+    r"""
+    decoder_sparse_step (`int`, *optional*, defaults to 1):
+        The frequency of the MoE layer.
+    mlp_only_layers (`List[int]`, *optional*, defaults to `[]`):
+        Indicate which layers use Qwen3VLMoeMLP rather than Qwen3VLMoeSparseMoeBlock
+        The list contains layer index, from 0 to num_layers-1 if we have num_layers layers
+        If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity.
+
+    ```python
+    >>> from transformers import Qwen3VLMoeForConditionalGeneration, Qwen3VLMoeConfig
+
+    >>> # Initializing a Qwen3VLMoe style configuration
+    >>> configuration = Qwen3VLMoeConfig()
+
+    >>> # Initializing a model from the Qwen3-VL-30B-A3B style configuration
+    >>> model = Qwen3VLMoeForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "qwen3_vl_moe_text"
+    base_config_key = "text_config"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    default_theta = 500000.0
+    # Default tensor parallel plan for base model `Qwen3VLMoe`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_ep_plan = {
+        "layers.*.mlp.gate": "ep_router",
+        "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+        "layers.*.mlp.experts.down_proj": "grouped_gemm",
+        "layers.*.mlp.experts": "moe_tp_experts",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+    ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"}
+
+    intermediate_size: int = 5632
+    num_hidden_layers: int = 24
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 16
+    max_position_embeddings: int = 128000
+    moe_intermediate_size: int = 1408
+    num_experts_per_tok: int = 4
+    num_experts: int = 60
+    head_dim: int | None = None
+    tie_word_embeddings: bool = True
+
+    norm_topk_prob = AttributeError()
+    output_router_logits = AttributeError()
+    use_sliding_window = AttributeError()
+    sliding_window = AttributeError()
+
+    def __post_init__(self, **kwargs):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+        self.head_dim = self.head_dim or self.hidden_size // self.num_attention_heads
+        super().__post_init__(**kwargs)
+        self.sliding_window = None
+
+
+@auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct")
+@strict
+class Qwen3VLMoeVisionConfig(Qwen3VLVisionConfig):
+    model_type = "qwen3_vl_moe_vision"
+    pass
+
+
+@auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct")
+@strict
+class Qwen3VLMoeConfig(Qwen3VLConfig):
+    r"""
+    Example:
+
+    ```python
+    >>> from transformers import Qwen3VLMoeForConditionalGeneration, Qwen3VLMoeConfig
+
+    >>> # Initializing a Qwen3-VL-MOE style configuration
+    >>> configuration = Qwen3VLMoeConfig()
+
+    >>> # Initializing a model from the Qwen3-VL-30B-A3B style configuration
+    >>> model = Qwen3VLMoeForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    pass
+
+
+class Qwen3VLMoeTextRMSNorm(Qwen3MoeRMSNorm):
+    pass
+
+
+class Qwen3VLMoeTextExperts(Qwen3MoeExperts):
+    pass
+
+
+class Qwen3VLMoeTextTopKRouter(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+        self.hidden_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
+
+    def forward(self, hidden_states):
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = F.linear(hidden_states, self.weight)  # (seq_len, num_experts)
+        router_probs = torch.nn.functional.softmax(router_logits, dtype=torch.float, dim=-1)
+        router_top_value, router_indices = torch.topk(router_probs, self.top_k, dim=-1)  # (seq_len, top_k)
+        router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
+        router_top_value = router_top_value.to(router_logits.dtype)
+        router_scores = router_top_value
+        return router_logits, router_scores, router_indices
+
+
+class Qwen3VLMoeTextSparseMoeBlock(Qwen3MoeSparseMoeBlock):
+    pass
+
+
+class Qwen3VLMoeTextAttention(Qwen3VLTextAttention):
+    pass
+
+
+class Qwen3VLMoeTextDecoderLayer(Qwen3MoeDecoderLayer):
+    pass
+
+
+class Qwen3VLMoePreTrainedModel(Qwen3MoePreTrainedModel):
+    config: Qwen3VLMoeConfig
+    input_modalities = ("text", "image", "video")
+    _no_split_modules = ["Qwen3VLMoeTextDecoderLayer", "Qwen3VLMoeVisionBlock"]
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        """Initialize the weights."""
+        PreTrainedModel._init_weights(self, module)
+        if hasattr(self.config, "initializer_range"):
+            std = self.config.initializer_range
+        else:
+            std = getattr(self.config.get_text_config(), "initializer_range", 0.02)
+        if isinstance(module, Qwen3VLMoeTextExperts):
+            init.normal_(module.gate_up_proj, mean=0.0, std=std)
+            init.normal_(module.down_proj, mean=0.0, std=std)
+        elif isinstance(module, Qwen3VLMoeTextTopKRouter):
+            init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, Qwen3VLMoeVisionRotaryEmbedding):
+            inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim))
+            init.copy_(module.inv_freq, inv_freq)
+
+
+class Qwen3VLMoeVisionRotaryEmbedding(Qwen3VLVisionRotaryEmbedding):
+    pass
+
+
+class Qwen3VLMoeVisionAttention(Qwen3VLVisionAttention):
+    pass
+
+
+class Qwen3VLMoeVisionBlock(Qwen3VLVisionBlock):
+    pass
+
+
+class Qwen3VLMoeVisionModel(Qwen3VLVisionModel):
+    _can_record_outputs = {
+        "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, layer_name="mlp.gate", index=0),
+        "hidden_states": Qwen3VLMoeVisionBlock,
+        "attentions": Qwen3VLMoeVisionAttention,
+    }
+
+
+class Qwen3VLMoeTextModel(Qwen3VLTextModel):
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        # args for deepstack
+        visual_pos_masks: torch.Tensor | None = None,
+        deepstack_visual_embeds: list[torch.Tensor] | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple | MoeModelOutputWithPast:
+        r"""
+        visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*):
+            The mask of the visual positions.
+        deepstack_visual_embeds (`list[torch.Tensor]`, *optional*):
+            The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim).
+            The feature is extracted from the different visual encoder layers, and fed to the decoder
+            hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334).
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        # torch.jit.trace() doesn't support cache objects in the output
+        if use_cache and past_key_values is None and not torch.jit.is_tracing():
+            past_key_values = DynamicCache(config=self.config)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # the hard coded `4` is for text, temporal, height and width.
+        if position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1)
+        elif position_ids.ndim == 2:
+            position_ids = position_ids[None, ...].expand(4, position_ids.shape[0], -1)
+
+        if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+            text_position_ids = position_ids[0]
+            position_ids = position_ids[1:]
+        else:
+            text_position_ids = None
+
+        attention_mask = create_causal_mask(
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_ids=text_position_ids,
+        )
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        for layer_idx, decoder_layer in enumerate(self.layers):
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=text_position_ids,
+                past_key_values=past_key_values,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+            hidden_states = layer_outputs
+
+            # add visual features to the hidden states of first several layers
+            if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
+                hidden_states = self._deepstack_process(
+                    hidden_states,
+                    visual_pos_masks,
+                    deepstack_visual_embeds[layer_idx],
+                )
+
+        hidden_states = self.norm(hidden_states)
+
+        return MoeModelOutputWithPast(  # only diff with Qwen3VLTextModel
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+class Qwen3VLMoeModelOutputWithPast(Qwen3VLModelOutputWithPast):
+    router_logits: tuple[torch.FloatTensor] | None = None
+
+
+class Qwen3VLMoeCausalLMOutputWithPast(Qwen3VLCausalLMOutputWithPast):
+    router_logits: tuple[torch.FloatTensor] | None = None
+    aux_loss: torch.FloatTensor | None = None
+
+
+class Qwen3VLMoeForConditionalGeneration(Qwen3VLForConditionalGeneration):
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        pixel_values: torch.Tensor | None = None,
+        pixel_values_videos: torch.FloatTensor | None = None,
+        image_grid_thw: torch.LongTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        mm_token_type_ids: torch.IntTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | Qwen3VLMoeCausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            The temporal, height and width of feature shape of each image in LLM.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            The temporal, height and width of feature shape of each video in LLM.
+
+        Example:
+        ```python
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+        >>> from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
+
+        >>> model = Qwen3VLMoeForConditionalGeneration.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto")
+        >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct")
+
+        >>> messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
+                    },
+                    {"type": "text", "text": "Describe this image in short."},
+                ],
+            }
+        ]
+
+        >>> # Preparation for inference
+        >>> inputs = processor.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt"
+        )
+        >>> inputs = inputs.to(model.device)
+
+        >>> # Generate
+        >>> generated_ids = model.generate(**inputs, max_new_tokens=128)
+        >>> generated_ids_trimmed = [
+            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        >>> processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "A woman in a plaid shirt sits on a sandy beach at sunset, smiling as she gives a high-five to a yellow Labrador Retriever wearing a harness. The ocean waves roll in the background."
+        ```"""
+
+        outputs = self.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            mm_token_type_ids=mm_token_type_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+
+        aux_loss = None
+        if kwargs.get("output_router_logits", False):
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits,
+                self.config.text_config.num_experts,
+                self.config.text_config.num_experts_per_tok,
+                attention_mask,
+            )
+            if labels is not None:
+                loss += self.config.text_config.router_aux_loss_coef * aux_loss.to(
+                    loss.device
+                )  # make sure to reside in the same device
+
+        return Qwen3VLMoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            rope_deltas=outputs.rope_deltas,
+            router_logits=outputs.router_logits,
+        )
+
+
+__all__ = [
+    "Qwen3VLMoeConfig",
+    "Qwen3VLMoeTextConfig",
+    "Qwen3VLMoeVisionConfig",
+    "Qwen3VLMoeVisionModel",
+    "Qwen3VLMoeForConditionalGeneration",
+    "Qwen3VLMoeModel",  # noqa
+    "Qwen3VLMoePreTrainedModel",
+    "Qwen3VLMoeTextModel",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/logs/infer_not5_bottleneck128_170k_decode32_ema_20260611/lr4e3.log

@@ -0,0 +1,29 @@
+checkpoint=runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_2x8gpu_50ep_lr4e3_ema0p9999_elfopt_not5_bottleneck128_unfixed_norm_stateprobadd_selfcond_ce_fast_20260610_020108/step_170000.pt
+use_ema=1
+step=170000
+decode_steps=32
+n=64 chunk_n=8 gpu=3
+out_base=/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/docs/lta_samples/metrics_20260611
+[2026-06-11T21:37:16+00:00] infer step=170000 decode=32 -> /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/docs/lta_samples/metrics_20260611/owt_t5_not5_bottleneck128_norm_stateprobadd_selfcond_ce_fast_lr4e3_ema0p9999_step170000_ema_sc1p0_decode32_n64
+[2026-06-11T21:37:16+00:00] run decode=32 chunk=0 n=8 seed=123
+[2026-06-11T21:37:23+00:00] done decode=32 chunk=0
+[2026-06-11T21:37:23+00:00] run decode=32 chunk=1 n=8 seed=124
+[2026-06-11T21:37:29+00:00] done decode=32 chunk=1
+[2026-06-11T21:37:29+00:00] run decode=32 chunk=2 n=8 seed=125
+[2026-06-11T21:37:36+00:00] done decode=32 chunk=2
+[2026-06-11T21:37:36+00:00] run decode=32 chunk=3 n=8 seed=126
+[2026-06-11T21:37:43+00:00] done decode=32 chunk=3
+[2026-06-11T21:37:43+00:00] run decode=32 chunk=4 n=8 seed=127
+[2026-06-11T21:37:50+00:00] done decode=32 chunk=4
+[2026-06-11T21:37:50+00:00] run decode=32 chunk=5 n=8 seed=128
+[2026-06-11T21:37:57+00:00] done decode=32 chunk=5
+[2026-06-11T21:37:57+00:00] run decode=32 chunk=6 n=8 seed=129
+[2026-06-11T21:38:04+00:00] done decode=32 chunk=6
+[2026-06-11T21:38:04+00:00] run decode=32 chunk=7 n=8 seed=130
+[2026-06-11T21:38:11+00:00] done decode=32 chunk=7
+merged 64 samples -> /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/docs/lta_samples/metrics_20260611/owt_t5_not5_bottleneck128_norm_stateprobadd_selfcond_ce_fast_lr4e3_ema0p9999_step170000_ema_sc1p0_decode32_n64/sc1p0/samples64.txt
+loading scorer /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard dtype=fp32 device=cuda
+run	kind	ppl	mean_entropy	distinct_1	distinct_2	top_token_mass	eos_rows	eos_total	ppl_tokens	t5_tokens	path
+sc1p0	raw_full	34.21054852666053	5.149774446140075	0.08880782318138622	0.522677457427078	0.02952086079519328	63	64	61342	65242	/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/docs/lta_samples/metrics_20260611/owt_t5_not5_bottleneck128_norm_stateprobadd_selfcond_ce_fast_lr4e3_ema0p9999_step170000_ema_sc1p0_decode32_n64/sc1p0
+sc1p0	pre_eos	41.24451616751232	5.1832901751118445	0.09159038087558696	0.5390526182646092	0.027779095321665163	0	0	57318	63249	/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/docs/lta_samples/metrics_20260611/owt_t5_not5_bottleneck128_norm_stateprobadd_selfcond_ce_fast_lr4e3_ema0p9999_step170000_ema_sc1p0_decode32_n64/sc1p0
+[2026-06-11T21:38:24+00:00] done