SpeedrunTrainLoop(
        study_name='speedrundit',
        experiment_name='exp004',
        base_dir='/opt/scratch',
        working_dir=PosixPath('/opt/scratch/runs/speedrundit/exp004'),
        doc='',
        step=SpeedrunTrainStep.Config(
        │   model=SpeedrunDiT.Config(
        │       input_size=16,
        │       in_channels=768,
        │       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.float64,
        │       reference_rope=True,
        │       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=True,
        │               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=ElementwiseLatentStats.Config(
        │       stats=HubFile.Config(
        │           repo_id='nyu-visionx/RAE-collections',
        │           filename='stats/dinov2/wReg_base/imagenet1k/stat.pt',
        │           revision='1be4f03273523431f099a934da4cf1940dc6039f',
        │           sha256='84ede66def5e6e3f25679334dc89cf63b12aacb99cbf0f5ae7ed4ad3187f7e59'
        │       ),
        │       eps=1e-05
        │   ),
        │   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='rae-dinov2-base-u8',
        │   │   autoencoder=RAE.Config(
        │   │       encoder=Encoder.Config(
        │   │           channels_hidden=768,
        │   │           num_layers=12,
        │   │           heads=12,
        │   │           expansion=4,
        │   │           patch_size=14,
        │   │           image_size=518,
        │   │           num_register_tokens=4,
        │   │           eps=1e-06,
        │   │           activation=<built-in function gelu>,
        │   │           checkpoint=HubFile.Config(
        │   │               repo_id='facebook/dinov2-with-registers-base',
        │   │               filename='model.safetensors',
        │   │               revision='a1d738ccfa7ae170945f210395d99dde8adb1805',
        │   │               sha256='7a6f7b3b9fa4b8732e707476a03cd6cdce210048582f21aafb7991c17d98e362'
        │   │           )
        │   │       ),
        │   │       decoder=Decoder.Config(
        │   │           channels_in=768,
        │   │           num_patches=256,
        │   │           channels_hidden=1_152,
        │   │           channels_hidden_mlp=4_096,
        │   │           num_layers=28,
        │   │           heads=16,
        │   │           patch_size=16,
        │   │           eps=1e-12,
        │   │           activation=<built-in function gelu>,
        │   │           checkpoint=HubFile.Config(
        │   │               repo_id='nyu-visionx/RAE-collections',
        │   │               filename='decoders/dinov2/wReg_base/ViTXL_n08/model.pt',
        │   │               revision='1be4f03273523431f099a934da4cf1940dc6039f',
        │   │               sha256='5fedf7c9660476a709e122cef18385c917532914a23f034388e1c1a52bde2be6'
        │   │           )
        │   │       ),
        │   │       image_size=256,
        │   │       encoder_image_size=224,
        │   │       pixel_mean=(0.485, 0.456, 0.406),
        │   │       pixel_std=(0.229, 0.224, 0.225),
        │   │       latent_norm=ElementwiseLatentStats.Config(
        │   │           stats=HubFile.Config(
        │   │               repo_id='nyu-visionx/RAE-collections',
        │   │               filename='stats/dinov2/wReg_base/imagenet1k/stat.pt',
        │   │               revision='1be4f03273523431f099a934da4cf1940dc6039f',
        │   │               sha256='84ede66def5e6e3f25679334dc89cf63b12aacb99cbf0f5ae7ed4ad3187f7e59'
        │   │           ),
        │   │           eps=1e-05
        │   │       )
        │   │   ),
        │   │   codec=ScalarTableCodec.Config(
        │   │       fit=LloydMaxFit.Config(max_iterations=500, tolerance=1e-07),
        │   │       groups=PerChannel.Config(),
        │   │       num_fit_images=4_096
        │   │   ),
        │   │   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/exp004'),
            working_dir=PosixPath('/opt/scratch/runs/speedrundit/exp004/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/exp004'),
            working_dir=PosixPath('/opt/scratch/runs/speedrundit/exp004/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}
        )
    )
