WorldModelLoop.Config(
        study_name='craftax-world-model',
        experiment_name='exp001',
        base_dir='/opt/scratch',
        working_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001'),
        doc='',
        step=WorldModelTrainStep.Config(
        │   model=WorldModel.Config(
        │   │   schema=FrameSchema(
        │   │   │   vocab_size=461,
        │   │   │   cell_slots=99,
        │   │   │   cell_fields=(
        │   │   │       CellField(name='block', offset=0, size=64, valid=37),
        │   │   │       CellField(name='item', offset=64, size=8, valid=6),
        │   │   │       CellField(
        │   │   │           name='visibility',
        │   │   │           offset=72,
        │   │   │           size=2,
        │   │   │           valid=2
        │   │   │       ),
        │   │   │       CellField(name='melee', offset=74, size=16, valid=9),
        │   │   │       CellField(name='passive', offset=90, size=16, valid=9),
        │   │   │       CellField(name='ranged', offset=106, size=16, valid=9),
        │   │   │       CellField(
        │   │   │           name='mob_projectile',
        │   │   │           offset=122,
        │   │   │           size=16,
        │   │   │           valid=9
        │   │   │       ),
        │   │   │       CellField(
        │   │   │           name='player_projectile',
        │   │   │           offset=138,
        │   │   │           size=16,
        │   │   │           valid=9
        │   │   │       )
        │   │   │   ),
        │   │   │   scalar_names=(
        │   │   │   │   'wood',
        │   │   │   │   'stone',
        │   │   │   │   'coal',
        │   │   │   │   'iron',
        │   │   │   │   'diamond',
        │   │   │   │   'sapphire',
        │   │   │   │   'ruby',
        │   │   │   │   'sapling',
        │   │   │   │   'torches',
        │   │   │   │   'arrows',
        │   │   │   │   'books',
        │   │   │   │   'pickaxe',
        │   │   │   │   'sword',
        │   │   │   │   'sword_enchantment',
        │   │   │   │   'bow_enchantment',
        │   │   │   │   'bow',
        │   │   │   │   'potion_red',
        │   │   │   │   'potion_green',
        │   │   │   │   'potion_blue',
        │   │   │   │   'potion_pink',
        │   │   │   │   'potion_cyan',
        │   │   │   │   'potion_yellow',
        │   │   │   │   'health',
        │   │   │   │   'food',
        │   │   │   │   'drink',
        │   │   │   │   'energy',
        │   │   │   │   'mana',
        │   │   │   │   'xp',
        │   │   │   │   'dexterity',
        │   │   │   │   'strength',
        │   │   │   │   'intelligence',
        │   │   │   │   'facing_left',
        │   │   │   │   'facing_right',
        │   │   │   │   'facing_up',
        │   │   │   │   'facing_down',
        │   │   │   │   'armour_0',
        │   │   │   │   'armour_1',
        │   │   │   │   'armour_2',
        │   │   │   │   'armour_3',
        │   │   │   │   'armour_enchantment_0',
        │   │   │   │   'armour_enchantment_1',
        │   │   │   │   'armour_enchantment_2',
        │   │   │   │   'armour_enchantment_3',
        │   │   │   │   'light',
        │   │   │   │   'sleeping',
        │   │   │   │   'resting',
        │   │   │   │   'learned_fireball',
        │   │   │   │   'learned_iceball',
        │   │   │   │   'floor',
        │   │   │   │   'floor_clear',
        │   │   │   │   'boss_vulnerable'
        │   │   │   ),
        │   │   │   scalar_ranges=(
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 159),
        │   │   │   │   (155, 159),
        │   │   │   │   (155, 157),
        │   │   │   │   (155, 157),
        │   │   │   │   (155, 156),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 254),
        │   │   │   │   (155, 415),
        │   │   │   │   (155, 172),
        │   │   │   │   (155, 172),
        │   │   │   │   (155, 172),
        │   │   │   │   (155, 176),
        │   │   │   │   (155, 163),
        │   │   │   │   (156, 160),
        │   │   │   │   (156, 160),
        │   │   │   │   (156, 160),
        │   │   │   │   (155, 156),
        │   │   │   │   (155, 156),
        │   │   │   │   (155, 156),
        │   │   │   │   (155, 156),
        │   │   │   │   (155, 157),
        │   │   │   │   (155, 157),
        │   │   │   │   (155, 157),
        │   │   │   │   (155, 157),
        │   │   │   │   (155, 157),
        │   │   │   │   (155, 157),
        │   │   │   │   (155, 157),
        │   │   │   │   (155, 157),
        │   │   │   │   (155, 410),
        │   │   │   │   (155, 156),
        │   │   │   │   (155, 156),
        │   │   │   │   (155, 156),
        │   │   │   │   (155, 156),
        │   │   │   │   (155, 163),
        │   │   │   │   (155, 156),
        │   │   │   │   (155, 156)
        │   │   │   ),
        │   │   │   prefix_names=('reward', 'done'),
        │   │   │   prefix_ranges=((154, 389), (459, 460))
        │   │   ),
        │   │   scalar_offset=155,
        │   │   num_actions=43,
        │   │   encoder=FrameEncoder.Config(
        │   │   │   channels_in=512,
        │   │   │   num_slots=150,
        │   │   │   stack=Transformer.Config(
        │   │   │   │   channels_in=512,
        │   │   │   │   channels_out=512,
        │   │   │   │   num_layers=4,
        │   │   │   │   proj_in=None,
        │   │   │   │   block=[
        │   │   │   │   │   EncoderBlock.Config(
        │   │   │   │   │       channels_in=512,
        │   │   │   │   │       channels_out=512,
        │   │   │   │   │       attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=False,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=((0, 4),),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           channels_hidden=1_344,
        │   │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │   │           round_to=256,
        │   │   │   │   │           gate=True,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           split_gate_projection=False,
        │   │   │   │   │           norm=None,
        │   │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │   │           depth_index=((0, 4),),
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           up_proj=None,
        │   │   │   │   │           middle=None,
        │   │   │   │   │           down_proj=None,
        │   │   │   │   │           shard='colwise'
        │   │   │   │   │       ),
        │   │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       prenorm=True,
        │   │   │   │   │       checkpoint=False,
        │   │   │   │   │       depth_index=((0, 4),),
        │   │   │   │   │       recompute_policy=None
        │   │   │   │   │   ),
        │   │   │   │   │   EncoderBlock.Config(
        │   │   │   │   │       channels_in=512,
        │   │   │   │   │       channels_out=512,
        │   │   │   │   │       attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=False,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=((1, 4),),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           channels_hidden=1_344,
        │   │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │   │           round_to=256,
        │   │   │   │   │           gate=True,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           split_gate_projection=False,
        │   │   │   │   │           norm=None,
        │   │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │   │           depth_index=((1, 4),),
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           up_proj=None,
        │   │   │   │   │           middle=None,
        │   │   │   │   │           down_proj=None,
        │   │   │   │   │           shard='colwise'
        │   │   │   │   │       ),
        │   │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       prenorm=True,
        │   │   │   │   │       checkpoint=False,
        │   │   │   │   │       depth_index=((1, 4),),
        │   │   │   │   │       recompute_policy=None
        │   │   │   │   │   ),
        │   │   │   │   │   EncoderBlock.Config(
        │   │   │   │   │       channels_in=512,
        │   │   │   │   │       channels_out=512,
        │   │   │   │   │       attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=False,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=((2, 4),),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           channels_hidden=1_344,
        │   │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │   │           round_to=256,
        │   │   │   │   │           gate=True,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           split_gate_projection=False,
        │   │   │   │   │           norm=None,
        │   │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │   │           depth_index=((2, 4),),
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           up_proj=None,
        │   │   │   │   │           middle=None,
        │   │   │   │   │           down_proj=None,
        │   │   │   │   │           shard='colwise'
        │   │   │   │   │       ),
        │   │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       prenorm=True,
        │   │   │   │   │       checkpoint=False,
        │   │   │   │   │       depth_index=((2, 4),),
        │   │   │   │   │       recompute_policy=None
        │   │   │   │   │   ),
        │   │   │   │   │   EncoderBlock.Config(
        │   │   │   │   │       channels_in=512,
        │   │   │   │   │       channels_out=512,
        │   │   │   │   │       attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=False,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=((3, 4),),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           channels_hidden=1_344,
        │   │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │   │           round_to=256,
        │   │   │   │   │           gate=True,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           split_gate_projection=False,
        │   │   │   │   │           norm=None,
        │   │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │   │           depth_index=((3, 4),),
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           up_proj=None,
        │   │   │   │   │           middle=None,
        │   │   │   │   │           down_proj=None,
        │   │   │   │   │           shard='colwise'
        │   │   │   │   │       ),
        │   │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       prenorm=True,
        │   │   │   │   │       checkpoint=False,
        │   │   │   │   │       depth_index=((3, 4),),
        │   │   │   │   │       recompute_policy=None
        │   │   │   │   │   )
        │   │   │   │   ],
        │   │   │   │   proj_out=RMSNorm.Config(
        │   │   │   │       channels_in=512,
        │   │   │   │       channels_out=512,
        │   │   │   │       eps=1e-06,
        │   │   │   │       elementwise_affine=True,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None
        │   │   │   │   )
        │   │   │   ),
        │   │   │   init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   cast_stream=None
        │   │   ),
        │   │   transformer=GlobalTransformer.Config(
        │   │   │   channels_in=1_152,
        │   │   │   channels_out=1_152,
        │   │   │   num_layers=20,
        │   │   │   proj_in=None,
        │   │   │   block=[
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((0, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((0, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((0, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((1, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((1, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((1, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((2, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((2, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((2, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((3, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((3, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((3, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((4, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((4, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((4, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((5, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((5, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((5, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((6, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((6, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((6, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((7, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((7, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((7, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((8, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((8, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((8, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((9, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((9, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((9, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((10, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((10, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((10, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((11, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((11, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((11, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((12, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((12, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((12, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((13, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((13, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((13, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((14, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((14, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((14, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((15, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((15, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((15, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((16, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((16, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((16, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((17, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((17, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((17, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((18, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((18, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((18, 20),)
        │   │   │   │   ),
        │   │   │   │   TransformerBlock.Config(
        │   │   │   │       channels_in=1_152,
        │   │   │   │       channels_out=1_152,
        │   │   │   │       attn=VarlenAttention.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           num_heads=9,
        │   │   │   │           channels_head=128,
        │   │   │   │           num_heads_kv=3,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.0,
        │   │   │   │           causal=True,
        │   │   │   │           rope=RoPE.Config(
        │   │   │   │               channels_head=128,
        │   │   │   │               reduction_mode='cat',
        │   │   │   │               frequencies=HuggingFaceFrequencies.Config(
        │   │   │   │                   base=500000.0
        │   │   │   │               ),
        │   │   │   │               dtype=torch.float32
        │   │   │   │           ),
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │               channels_in=128,
        │   │   │   │               channels_out=128,
        │   │   │   │               eps=1e-06,
        │   │   │   │               elementwise_affine=True,
        │   │   │   │               device=None,
        │   │   │   │               dtype=None
        │   │   │   │           ),
        │   │   │   │           share_qk_norm=False,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=None,
        │   │   │   │           proj_out=None,
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           depth_index=((19, 20),),
        │   │   │   │           attn_kernel=Flash4Varlen.Config()
        │   │   │   │       ),
        │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           channels_hidden=3_072,
        │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │           round_to=256,
        │   │   │   │           gate=True,
        │   │   │   │           bias=False,
        │   │   │   │           split_gate_projection=False,
        │   │   │   │           norm=None,
        │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │           depth_index=((19, 20),),
        │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │           up_proj=None,
        │   │   │   │           middle=None,
        │   │   │   │           down_proj=None,
        │   │   │   │           shard='colwise'
        │   │   │   │       ),
        │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │           channels_in=1_152,
        │   │   │   │           channels_out=1_152,
        │   │   │   │           eps=1e-06,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       prenorm=True,
        │   │   │   │       checkpoint=False,
        │   │   │   │       depth_index=((19, 20),)
        │   │   │   │   )
        │   │   │   ],
        │   │   │   proj_out=RMSNorm.Config(
        │   │   │       channels_in=1_152,
        │   │   │       channels_out=1_152,
        │   │   │       eps=1e-06,
        │   │   │       elementwise_affine=True,
        │   │   │       device=None,
        │   │   │       dtype=None
        │   │   │   )
        │   │   ),
        │   │   decoder=LocalDecoder.Config(
        │   │   │   channels_in=512,
        │   │   │   num_slots=152,
        │   │   │   stack=Transformer.Config(
        │   │   │   │   channels_in=512,
        │   │   │   │   channels_out=512,
        │   │   │   │   num_layers=4,
        │   │   │   │   proj_in=None,
        │   │   │   │   block=[
        │   │   │   │   │   DecoderBlock.Config(
        │   │   │   │   │       channels_in=512,
        │   │   │   │   │       channels_out=512,
        │   │   │   │   │       attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=True,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       cross_attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=False,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           channels_hidden=1_344,
        │   │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │   │           round_to=256,
        │   │   │   │   │           gate=True,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           split_gate_projection=False,
        │   │   │   │   │           norm=None,
        │   │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           up_proj=None,
        │   │   │   │   │           middle=None,
        │   │   │   │   │           down_proj=None,
        │   │   │   │   │           shard=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm3=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       checkpoint=True,
        │   │   │   │   │       recompute_policy=None
        │   │   │   │   │   ),
        │   │   │   │   │   DecoderBlock.Config(
        │   │   │   │   │       channels_in=512,
        │   │   │   │   │       channels_out=512,
        │   │   │   │   │       attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=True,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       cross_attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=False,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           channels_hidden=1_344,
        │   │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │   │           round_to=256,
        │   │   │   │   │           gate=True,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           split_gate_projection=False,
        │   │   │   │   │           norm=None,
        │   │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           up_proj=None,
        │   │   │   │   │           middle=None,
        │   │   │   │   │           down_proj=None,
        │   │   │   │   │           shard=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm3=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       checkpoint=True,
        │   │   │   │   │       recompute_policy=None
        │   │   │   │   │   ),
        │   │   │   │   │   DecoderBlock.Config(
        │   │   │   │   │       channels_in=512,
        │   │   │   │   │       channels_out=512,
        │   │   │   │   │       attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=True,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       cross_attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=False,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           channels_hidden=1_344,
        │   │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │   │           round_to=256,
        │   │   │   │   │           gate=True,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           split_gate_projection=False,
        │   │   │   │   │           norm=None,
        │   │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           up_proj=None,
        │   │   │   │   │           middle=None,
        │   │   │   │   │           down_proj=None,
        │   │   │   │   │           shard=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm3=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       checkpoint=True,
        │   │   │   │   │       recompute_policy=None
        │   │   │   │   │   ),
        │   │   │   │   │   DecoderBlock.Config(
        │   │   │   │   │       channels_in=512,
        │   │   │   │   │       channels_out=512,
        │   │   │   │   │       attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=True,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       cross_attn=Attention.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=64,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.0,
        │   │   │   │   │           causal=False,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=RMSNorm.Config(
        │   │   │   │   │               channels_in=64,
        │   │   │   │   │               channels_out=64,
        │   │   │   │   │               eps=1e-06,
        │   │   │   │   │               elementwise_affine=True,
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None
        │   │   │   │   │           ),
        │   │   │   │   │           share_qk_norm=False,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=None,
        │   │   │   │   │           proj_out=None,
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           attn_kernel=SdpaFused.Config()
        │   │   │   │   │       ),
        │   │   │   │   │       ffn=SwiGLU.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           channels_hidden=1_344,
        │   │   │   │   │           expansion=2.6666666666666665,
        │   │   │   │   │           round_to=256,
        │   │   │   │   │           gate=True,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           split_gate_projection=False,
        │   │   │   │   │           norm=None,
        │   │   │   │   │           act=priml.model.swiglu.silu,
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   │   │           up_proj=None,
        │   │   │   │   │           middle=None,
        │   │   │   │   │           down_proj=None,
        │   │   │   │   │           shard=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm1=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm2=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       norm3=RMSNorm.Config(
        │   │   │   │   │           channels_in=512,
        │   │   │   │   │           channels_out=512,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       checkpoint=True,
        │   │   │   │   │       recompute_policy=None
        │   │   │   │   │   )
        │   │   │   │   ],
        │   │   │   │   proj_out=RMSNorm.Config(
        │   │   │   │       channels_in=512,
        │   │   │   │       channels_out=512,
        │   │   │   │       eps=1e-06,
        │   │   │   │       elementwise_affine=True,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None
        │   │   │   │   )
        │   │   │   ),
        │   │   │   init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   │   cast_stream=None
        │   │   ),
        │   │   init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │   │   z_loss=1e-05,
        │   │   output_table=None,
        │   │   embed_board=priml.baselines.craftax.world_model.model.gathered_board
        │   ),
        │   optimizer=CompositeOptimizer.Config(
        │       optimizers=[
        │           FusedAdamW.Config(
        │               lr=0.0003,
        │               betas=(0.9, 0.95),
        │               eps=1e-08,
        │               weight_decay=0.0,
        │               compile=True
        │           ),
        │           Muon.Config(
        │               lr=0.0003,
        │               weight_decay=0.1,
        │               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_match_rms_adamw,
        │               momentum_update_fn=None,
        │               matrix_signum_fn=None,
        │               ensemble_dims=0
        │           ),
        │           Muon.Config(
        │               lr=0.0003,
        │               weight_decay=0.1,
        │               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_match_rms_adamw,
        │               momentum_update_fn=None,
        │               matrix_signum_fn=None,
        │               ensemble_dims=1
        │           )
        │       ],
        │       select=[
        │           excluding(complement(excluding(Muon.eligible_tensor, 'embedding', 'table', 'action_head', 'proj_qkv')), 'proj_qkv'),
        │           excluding(Muon.eligible_tensor, 'embedding', 'table', 'action_head', 'proj_qkv'),
        │           matching('proj_qkv')
        │       ],
        │       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=8,
        │   drop_partial_accumulation_on_epoch_end=True,
        │   train_budget_steps=1_525,
        │   train_budget_sec=inf,
        │   train_budget_epochs=inf,
        │   lr_schedule=PartialConfig(
        │       priml.baselines.craftax.world_model.train_step.wsd,
        │       decay=0.1,
        │       warmup=0.1
        │   ),
        │   parallelism=NoParallel.Config(device=None),
        │   model_quantization=NoModelQuantization.Config(),
        │   activation_memoization=DefaultActivationStorage.Config(),
        │   compile=PartialConfig(torch.compile, fullgraph=True),
        │   ema=NoEMA.Config(),
        │   gradient_clip_norm=1.0,
        │   skip_step_on_nonfinite_grad=False,
        │   device_init='eager',
        │   dtype_autocast=torch.bfloat16,
        │   autocast_cache_enabled=False,
        │   target_nll_fn=PartialConfig(
        │       priml.baselines.craftax.world_model.metric.craftax_target_nll
        │   ),
        │   cost_fn=PartialConfig(
        │       priml.baselines.craftax.world_model.train_step.world_model_cost
        │   ),
        │   device_peak_flops=989000000000000.0,
        │   loss_smoothing=0.9
        ),
        dataset=ReplayStream.Config(
            base_dir='/opt/scratch',
            working_dir=PosixPath('/opt/scratch/datasets/craftax/world-model/archive-v1'),
            corpus=PosixPath('/opt/scratch/datasets/craftax/world-model/archive-v1/corpora/base.json'),
            index_dir=PosixPath('/opt/scratch/datasets/craftax/world-model/archive-v1/index'),
            windows=1,
            t_g=8_192,
            s_max=64,
            micro_batches_per_step=8,
            sampler_seed=0,
            stratum_power=0.5,
            cached_decisions=1_000_000,
            replay_block_decisions=4_096,
            replay_cached_episodes=1_024,
            validation=StratifiedWindows.Config(batches=64),
            train_spans=None,
            count_multiple=1,
            device='cuda'
        ),
        metrics_train={},
        metrics_eval={   'val': CraftaxBitsPerByte.Config(
        │       schema=FrameSchema(
        │           vocab_size=461,
        │           cell_slots=99,
        │           cell_fields=(
        │               CellField(name='block', offset=0, size=64, valid=37),
        │               CellField(name='item', offset=64, size=8, valid=6),
        │               CellField(
        │                   name='visibility',
        │                   offset=72,
        │                   size=2,
        │                   valid=2
        │               ),
        │               CellField(name='melee', offset=74, size=16, valid=9),
        │               CellField(name='passive', offset=90, size=16, valid=9),
        │               CellField(name='ranged', offset=106, size=16, valid=9),
        │               CellField(
        │                   name='mob_projectile',
        │                   offset=122,
        │                   size=16,
        │                   valid=9
        │               ),
        │               CellField(
        │                   name='player_projectile',
        │                   offset=138,
        │                   size=16,
        │                   valid=9
        │               )
        │           ),
        │           scalar_names=(
        │           │   'wood',
        │           │   'stone',
        │           │   'coal',
        │           │   'iron',
        │           │   'diamond',
        │           │   'sapphire',
        │           │   'ruby',
        │           │   'sapling',
        │           │   'torches',
        │           │   'arrows',
        │           │   'books',
        │           │   'pickaxe',
        │           │   'sword',
        │           │   'sword_enchantment',
        │           │   'bow_enchantment',
        │           │   'bow',
        │           │   'potion_red',
        │           │   'potion_green',
        │           │   'potion_blue',
        │           │   'potion_pink',
        │           │   'potion_cyan',
        │           │   'potion_yellow',
        │           │   'health',
        │           │   'food',
        │           │   'drink',
        │           │   'energy',
        │           │   'mana',
        │           │   'xp',
        │           │   'dexterity',
        │           │   'strength',
        │           │   'intelligence',
        │           │   'facing_left',
        │           │   'facing_right',
        │           │   'facing_up',
        │           │   'facing_down',
        │           │   'armour_0',
        │           │   'armour_1',
        │           │   'armour_2',
        │           │   'armour_3',
        │           │   'armour_enchantment_0',
        │           │   'armour_enchantment_1',
        │           │   'armour_enchantment_2',
        │           │   'armour_enchantment_3',
        │           │   'light',
        │           │   'sleeping',
        │           │   'resting',
        │           │   'learned_fireball',
        │           │   'learned_iceball',
        │           │   'floor',
        │           │   'floor_clear',
        │           │   'boss_vulnerable'
        │           ),
        │           scalar_ranges=(
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 159),
        │           │   (155, 159),
        │           │   (155, 157),
        │           │   (155, 157),
        │           │   (155, 156),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 254),
        │           │   (155, 415),
        │           │   (155, 172),
        │           │   (155, 172),
        │           │   (155, 172),
        │           │   (155, 176),
        │           │   (155, 163),
        │           │   (156, 160),
        │           │   (156, 160),
        │           │   (156, 160),
        │           │   (155, 156),
        │           │   (155, 156),
        │           │   (155, 156),
        │           │   (155, 156),
        │           │   (155, 157),
        │           │   (155, 157),
        │           │   (155, 157),
        │           │   (155, 157),
        │           │   (155, 157),
        │           │   (155, 157),
        │           │   (155, 157),
        │           │   (155, 157),
        │           │   (155, 410),
        │           │   (155, 156),
        │           │   (155, 156),
        │           │   (155, 156),
        │           │   (155, 156),
        │           │   (155, 163),
        │           │   (155, 156),
        │           │   (155, 156)
        │           ),
        │           prefix_names=('reward', 'done'),
        │           prefix_ranges=((154, 389), (459, 460))
        │       )
            )},
        checkpointer=Checkpointer.Config(
            base_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001'),
            working_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001/checkpoints'),
            filename='step_{step:08d}.pt',
            save_every=1_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/craftax-world-model/exp001'),
            working_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001/profiling'),
            heartbeat_interval_sec=20.0,
            fault_dump_interval_sec=40.0,
            cuda_events=False
        ),
        tracker=NanochatSeries.Config(
            tracker=TrackerList.Config(
                trackers={   'metrics': FileTracker.Config(
                        base_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001'),
                        working_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001/metrics.json'),
                        capture_prefix='eval/'
                    ),
                    'val': FileTracker.Config(
                        base_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001'),
                        working_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001/val.json'),
                        capture_prefix='val/'
                    ),
                    'wandb': AsyncTracker.Config(
                        tracker=WandbTracker.Config(
                            project='craftax-world-model',
                            name='exp001',
                            run_id='',
                            group='exp001',
                            mode='online',
                            base_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001'),
                            working_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001/wandb'),
                            capture_console=True,
                            replay_startup_logs=False,
                            allow_startup_failure=True,
                            ingestion=WandbIngestion(
                                init_timeout_sec=30.0,
                                service_wait_sec=30.0,
                                flush_interval_sec=15.0,
                                system_metrics=True,
                                system_metrics_interval_sec=60.0
                            ),
                            metric_step_metrics={},
                            run_config={},
                            notes=''
                        ),
                        enabled=True
                    )},
                base_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001'),
                working_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001')
            ),
            base_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001'),
            working_dir=PosixPath('/opt/scratch/runs/craftax-world-model/exp001')
        ),
        max_steps=1_525,
        max_epochs=inf,
        max_time=inf,
        max_time_kind='wall',
        max_eval_time=inf,
        eval_stop_on_time_limit=False,
        num_steps_eval=250,
        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=None,
        model_torch_seed=None,
        num_steps_garbage_collect=inf,
        mesh_dim_model_seed='pp',
        mesh_dim_data_seed='dp',
        runtime=SingleProcess.Config(
            device='cuda',
            deterministic=False,
            float32_matmul_precision=None
        )
    )
