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 'functools.partial'>, <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=<function priml.optimizers.muon.adjust_lr_original at 0xdefacedeface>,
                │                               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=<function priml.loss.simple_loss.bce_with_logits at 0xdefacedeface>,
                │               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(<class 'functools.partial'>, <function priml.math.schedules.constant at 0xdefacedeface>),
                │       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='module',
                │               warmup_seed=False,
                │               decay_schedule=<function priml.train.ema.constant_decay at 0xdefacedeface>,
                │               select=<function priml.optimizers.parameter_filter.trainable at 0xdefacedeface>
                │       ),
                │       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(
                │               variant='dinov2_vitb14',
                │               layer_indices=(12, 12, 12),
                │               image_size=256
                │       ),
                │       latent_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')
                        ),
                        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,
                num_steps_garbage_collect=inf,
                mesh_dim_model_seed='pp',
                mesh_dim_data_seed='dp',
                runtime=MultiProcess.Config(
                        device='auto',
                        deterministic=False,
                        float32_matmul_precision='high',
                        backend=None,
                        mesh_topology={'dp': -1, 'pp': 1, 'tp': 1}
                )
        )
