SpeedrunTrainLoop(
        study_name='speedrundit',
        experiment_name='exp001',
        base_dir='/opt/scratch',
        working_dir=PosixPath('/opt/scratch/runs/speedrundit/exp001'),
        doc='',
        step=SpeedrunTrainStep.Config(
        │   model=SpeedrunDiT.Config(
        │       input_size=16,
        │       in_channels=32,
        │       patch_size=1,
        │       hidden_size=768,
        │       depth=12,
        │       num_heads=12,
        │       num_classes=1_000,
        │       cls_channels=768,
        │       timestep_frequencies=256,
        │       projector_hidden=2_048,
        │       projection_depths=(2, 4, 6),
        │       mlp_ratio_min=2.0,
        │       mlp_ratio_max=6.0,
        │       drop_ratio=0.75,
        │       path_drop_prob=0.05,
        │       class_dropout_prob=0.1,
        │       qk_norm=True,
        │       position_compute_dtype=torch.float32,
        │       reference_rope=False,
        │       encoder_blocks=2,
        │       decoder_blocks=2
        │   ),
        │   optimizer=CompositeOptimizer.Config(
        │       optimizers=[
        │           PartialConfig(
        │               <class 'torch.optim.adamw.AdamW'>,
        │               lr=0.0001,
        │               betas=(0.9, 0.999),
        │               weight_decay=0.0,
        │               eps=1e-15
        │           ),
        │           Muon.Config(
        │               lr=0.001,
        │               weight_decay=0.0,
        │               momentum=0.95,
        │               nesterov=True,
        │               ns_coefficients=(3.4445, -4.775, 2.0315),
        │               eps=1e-07,
        │               ns_steps=5,
        │               reference_numerics=False,
        │               adjust_lr_fn=priml.optimizers.muon.adjust_lr_original,
        │               momentum_update_fn=None,
        │               matrix_signum_fn=None,
        │               ensemble_dims=0
        │           )
        │       ],
        │       select=[
        │           complement(excluding(Muon.eligible_tensor, 'x_embedder', 't_embedder', 'y_embedder', 'final_layer', 'cls_projector', 'adaLN_modulation')),
        │           excluding(Muon.eligible_tensor, 'x_embedder', 't_embedder', 'y_embedder', 'final_layer', 'cls_projector', 'adaLN_modulation')
        │       ],
        │       require_total=True,
        │       drop_empty=False
        │   ),
        │   loss=SimpleLoss.Config(
        │       loss_fn=priml.loss.simple_loss.bce_with_logits,
        │       target_key='label',
        │       kwargs={'reduction': 'none'},
        │       channels_out=-1
        │   ),
        │   accumulate_grad_batches=1,
        │   drop_partial_accumulation_on_epoch_end=True,
        │   train_budget_steps=400_000,
        │   train_budget_sec=inf,
        │   train_budget_epochs=inf,
        │   lr_schedule=PartialConfig(priml.math.schedules.constant),
        │   parallelism=DataParallel.Config(
        │       mesh_dim='dp',
        │       bucket_cap_mb=25,
        │       find_unused_parameters=False,
        │       gradient_as_bucket_view=False
        │   ),
        │   model_quantization=NoModelQuantization.Config(),
        │   activation_memoization=DefaultActivationStorage.Config(),
        │   compile=None,
        │   ema=EMA.Config(
        │       decay=0.9999,
        │       update_after_step=0,
        │       update_every=1,
        │       track_buffers=False,
        │       shadow_kind='param_dict',
        │       warmup_seed=False,
        │       decay_schedule=priml.train.ema.constant_decay,
        │       select=priml.optimizers.parameter_filter.trainable,
        │       update_fn=None,
        │       buffer_update_fn=priml.train.ema.copy_buffers
        │   ),
        │   gradient_clip_norm=1.0,
        │   skip_step_on_nonfinite_grad=False,
        │   device_init='eager',
        │   dtype_autocast=torch.bfloat16,
        │   autocast_cache_enabled=False,
        │   teacher=DinoV2Teacher.Config(
        │       repository='facebookresearch/dinov2:7764ea0f912e53c92e82eb78a2a1631e92725fc8',
        │       variant='dinov2_vitb14',
        │       weights_file='dinov2_vitb14_pretrain.pth',
        │       weights_sha256='0b8b82f85de91b424aded121c7e1dcc2b7bc6d0adeea651bf73a13307fad8c73',
        │       layer_indices=(12, 12, 12),
        │       image_size=256
        │   ),
        │   latent_norm=ScaleLatents.Config(scale=0.3099),
        │   projection_coeff=0.5,
        │   cls_coeff=0.03,
        │   cfm_coeff=0.05,
        │   cfm_weighting='uniform',
        │   time_shifting=True,
        │   shift_base=4_096
        ),
        dataset=SpeedrunImageNetData.Config(
        │   source=PairedImageLatentDataset.Config(
        │       base_dir=PosixPath('/opt/scratch'),
        │       working_dir=PosixPath('/opt/scratch/datasets/speedrundit'),
        │       latent_subdir='vae-in',
        │       autoencoder=INVAE.Config(
        │           encoder=Encoder.Config(
        │               channels_in=3,
        │               channels_hidden=(128, 128, 256, 256, 512),
        │               blocks_per_stage=2,
        │               attention=(False, False, False, False, True),
        │               channels_latent=32,
        │               num_groups=32,
        │               eps=1e-06,
        │               checkpoint=HubFile.Config(
        │                   repo_id='REPA-E/e2e-invae-hf',
        │                   filename='diffusion_pytorch_model.safetensors',
        │                   revision='523ddf74b061c01ed0aebb7904c6bdb2f33e06d0',
        │                   sha256='6b1f754e414a841893fd242ca233c5aeec4bbdb7d2eeab2f7956adf48742e500'
        │               )
        │           ),
        │           decoder=Decoder.Config(
        │               channels_out=3,
        │               channels_hidden=(128, 128, 256, 256, 512),
        │               blocks_per_stage=3,
        │               attention=(False, False, False, False, True),
        │               channels_latent=32,
        │               num_groups=32,
        │               eps=1e-06,
        │               checkpoint=HubFile.Config(
        │                   repo_id='REPA-E/e2e-invae-hf',
        │                   filename='diffusion_pytorch_model.safetensors',
        │                   revision='523ddf74b061c01ed0aebb7904c6bdb2f33e06d0',
        │                   sha256='6b1f754e414a841893fd242ca233c5aeec4bbdb7d2eeab2f7956adf48742e500'
        │               )
        │           ),
        │           image_size=256,
        │           latent_norm=ScaleLatents.Config(scale=0.3099)
        │       ),
        │       codec=FloatCodec.Config(dtype=torch.float32),
        │       sample=True,
        │       seed=None
        │   ),
        │   batch_size=32,
        │   num_workers=4,
        │   prefetch_factor=2,
        │   pin_memory=True,
        │   base_dir='/opt/scratch',
        │   working_dir=PosixPath('/opt/scratch')
        ),
        metrics_train={},
        metrics_eval={},
        checkpointer=Checkpointer.Config(
            base_dir=PosixPath('/opt/scratch/runs/speedrundit/exp001'),
            working_dir=PosixPath('/opt/scratch/runs/speedrundit/exp001/checkpoints'),
            filename='step_{step:08d}.pt',
            save_every=10_000,
            keep_last_n=-1,
            keep_every=0,
            storer=SyncLocalStateDictStorer.Config(),
            resume=True,
            resume_step=-1,
            allow_checkpoint_overwrite=False,
            best_metric='',
            best_mode='max'
        ),
        profiler=None,
        phase_timer=PhaseTimer.Config(
            enabled=False,
            torch_profile=False,
            base_dir=PosixPath('/opt/scratch/runs/speedrundit/exp001'),
            working_dir=PosixPath('/opt/scratch/runs/speedrundit/exp001/profiling'),
            heartbeat_interval_sec=20.0,
            fault_dump_interval_sec=40.0,
            cuda_events=False
        ),
        tracker=None,
        max_steps=400_000,
        max_epochs=inf,
        max_time=inf,
        max_time_kind='wall',
        max_eval_time=inf,
        eval_stop_on_time_limit=False,
        num_steps_eval=inf,
        num_steps_log=10,
        early_train_log_steps=100,
        eval_every_epoch=False,
        eval_extras_every_eval=False,
        eval_warmup_batches=0,
        eval_only=False,
        restore_rng_state=True,
        seed=0,
        model_torch_seed=None,
        num_steps_garbage_collect=inf,
        mesh_dim_model_seed='pp',
        mesh_dim_data_seed='dp',
        runtime=MultiProcess.Config(
            device=None,
            deterministic=False,
            float32_matmul_precision='high',
            backend=None,
            process_group_timeout_sec=None,
            mesh_topology={'dp': -1, 'pp': 1, 'tp': 1}
        )
    )
