# Created by BVH, Jul - Aug 2025. # Converts the checkpoint from Cosmos to Any4D to accommodate and correctly populate new weights, depending on the config. import copy import math import torch from rich import print from custom.any4d.a4d_config import Any4DConfig from custom.any4d.a4d_network import Any4DDiT from imaginaire.utils import log def kaiming_uniform_like(tensor: torch.Tensor, fan_in_total: int = None): ''' Apply Kaiming uniform initialization but take total fan_in into account. ''' if fan_in_total is None: fan_in_total = tensor.shape[1] result = torch.rand_like(tensor) # values in [0, 1) bound = 1.0 / math.sqrt(fan_in_total) result = (result * 2.0 - 1.0) * bound # log.info(f'result: {result}') return result def validate_config(any4d_config: Any4DConfig): ''' Prepares the surgery by checking the config consistency and assigning useful helper variables. ''' if any4d_config.validated: log.debug('any4d_config already validated, skipping') return any4d_config num_views = len(any4d_config.video_entries) assert any4d_config.num_views == num_views # Determine actual channel ranges num_video_in_channels = [] num_video_out_channels = [] for view_idx in range(num_views): cur_channels = max([x[1] for x in any4d_config.video_entries[view_idx].values()]) num_video_in_channels.append(cur_channels) num_video_out_channels.append(cur_channels) # Same for simplicity. # Assign useful helper variables for network, VAE, etc any4d_config.num_video_in_channels = num_video_in_channels any4d_config.num_video_out_channels = num_video_out_channels any4d_config.num_lowdim_sattn_tokens = max([x[1] for x in any4d_config.lowdim_sattn_entries.values()]) \ if len(any4d_config.lowdim_sattn_entries) > 0 else 0 any4d_config.num_lowdim_sattn_channels = max([x[3] for x in any4d_config.lowdim_sattn_entries.values()]) \ if len(any4d_config.lowdim_sattn_entries) > 0 else 0 any4d_config.num_lowdim_xattn_tokens = max([x[1] for x in any4d_config.lowdim_xattn_entries.values()]) \ if len(any4d_config.lowdim_xattn_entries) > 0 else 0 any4d_config.num_lowdim_xattn_channels = max([x[3] for x in any4d_config.lowdim_xattn_entries.values()]) \ if len(any4d_config.lowdim_xattn_entries) > 0 else 0 any4d_config.num_lowdim_adaln_tokens = max([x[1] for x in any4d_config.lowdim_adaln_entries.values()]) \ if len(any4d_config.lowdim_adaln_entries) > 0 else 0 any4d_config.num_lowdim_adaln_channels = max([x[3] for x in any4d_config.lowdim_adaln_entries.values()]) \ if len(any4d_config.lowdim_adaln_entries) > 0 else 0 any4d_config.has_sattn_stream = (any4d_config.num_lowdim_sattn_tokens > 0) any4d_config.has_xattn_stream = (any4d_config.num_lowdim_xattn_tokens > 0) any4d_config.has_adaln_stream = (any4d_config.num_lowdim_adaln_tokens > 0) if any4d_config.any4d_active: any4d_config.all_highdim_entries = [k for v in any4d_config.video_entries for k in v.keys()] any4d_config.all_lowdim_entries = [k for k in any4d_config.lowdim_sattn_entries.keys()] any4d_config.all_lowdim_entries += [k for k in any4d_config.lowdim_xattn_entries.keys()] any4d_config.all_lowdim_entries += [k for k in any4d_config.lowdim_adaln_entries.keys()] else: # Vidar / vanilla fallback: any4d_config.all_highdim_entries = ['rgb0', 'depth0', 'points0', 'cams0'] any4d_config.all_lowdim_entries = ['action', 'proprio'] # Run sanity checks to catch mistakes early # TODO(bvh) # Verify each key (entry name) is unique, across all streams and views # Verify each key has associated one-channel input and output masks # Verify no unassigned / missing indices, and no overlapping / duplicate indices # Verify intactness of base model / pretrained functionality assert 'rgb0' in any4d_config.video_entries[0] assert 'rgb0_input_mask' in any4d_config.video_entries[0] assert tuple(any4d_config.video_entries[0]['rgb0']) == (0, 16, 'load', 'load') assert tuple(any4d_config.video_entries[0]['rgb0_input_mask']) == (16, 17, 'load', None) any4d_config.validated = True return any4d_config def load_checkpoint(dit: Any4DDiT, ckpt: dict, any4d_config: Any4DConfig): ''' Performs the necessary tensor logistics to map checkpoints from disk to actual architectures. :param dit (Any4DDiT): the network to load the checkpoint into. :param ckpt (dict): state_dict_dit_compatible, as loaded from dit_path, without net. prefix. :param any4d_config (Any4DConfig): the config to use to direct the surgery. ''' # DEBUG: if 0: (mk, uk) = dit.load_state_dict(ckpt, strict=False, assign=True) sd2 = {k: v for (k, v) in ckpt.items() if not('blocks' in k)} # for vanilla v2w: # sd2 = { # 'accum_video_sample_counter': tensor i64 4575657222473777152, # 'accum_image_sample_counter': tensor i64 4575657222473777152, # 'accum_iteration': tensor i64 4575657222473777152, # 'accum_train_in_hours': tensor bf16 1.000, # 'x_embedder.proj.1.weight': tensor[2048, 72] bf16 n=147456 (0.3Mb) x∈[-0.213, 0.146] μ=-4.917e-06 σ=0.011, # 'pos_embedder.seq': tensor[128] bf16 x∈[0., 127.000] μ=63.500 σ=37.000, # 'pos_embedder.dim_spatial_range': tensor[21] bf16 x∈[0., 0.953] μ=0.477 σ=0.295, # 'pos_embedder.dim_temporal_range': tensor[22] bf16 x∈[0., 0.953] μ=0.477 σ=0.295, # 't_embedder.1.linear_1.weight': tensor[2048, 2048] bf16 n=4194304 (8Mb) x∈[-0.189, 0.033] μ=-0.001 σ=0.006, # 't_embedder.1.linear_2.weight': tensor[6144, 2048] bf16 n=12582912 (24Mb) x∈[-0.188, 0.254] μ=-0.001 σ=0.007, # 'final_layer.linear.weight': tensor[64, 2048] bf16 n=131072 (0.2Mb) x∈[-0.046, 0.052] μ=-3.025e-06 σ=0.004, # 'final_layer.adaln_modulation.1.weight': tensor[256, 2048] bf16 n=524288 (1Mb) x∈[-0.098, 0.092] μ=0.000 σ=0.022, # 'final_layer.adaln_modulation.2.weight': tensor[4096, 256] bf16 n=1048576 (2Mb) x∈[-0.344, 0.369] μ=0.000 σ=0.026, # 't_embedding_norm.weight': tensor[2048] bf16 4Kb x∈[0.033, 0.072] μ=0.055 σ=0.010, # 't_embedding_norm._extra_state': tensor[5] u8 x∈[4, 148] μ=90.200 σ=61.977 [128, 4, 125, 148, 46] # } # ========================================================== # ======== Handle high-dimensional inputs & outputs ======== # ========================================================== V = any4d_config.num_views D = any4d_config.pipe_config.state_ch # 16 PS = any4d_config.pipe_config.net.patch_spatial # 2 PT = any4d_config.pipe_config.net.patch_temporal # 1 PP = PS * PS * PT # 4 ckpt_src = {k: v for (k, v) in ckpt.items() if not('blocks' in k)} ckpt_dst = copy.deepcopy(ckpt_src) # NOTE(bvh): There are no biases (on both sides); # only randomly initialized view embeddings that get summed as needed. # NOTE(bvh): For output projection, we intentionally ignore adaln_modulation # as it is always zero-initialized by default. in_prefixes = ['x_embedder.proj.1.weight'] + \ [f'x_embedder_newviews.{view_idx - 1}.proj.1.weight' for view_idx in range(1, V)] out_prefixes = ['final_layer.linear.weight'] + \ [f'final_layer_newviews.{view_idx - 1}.linear.weight' for view_idx in range(1, V)] key2view = {key: v for v in range(V) for key in any4d_config.video_entries[v].keys()} # Existing / root / pretrained viewpoint weights: base_in_prefix = in_prefixes[0] base_out_prefix = out_prefixes[0] in_embdim = ckpt_src[base_in_prefix].shape[0] # 2048 out_embdim = ckpt_src[base_out_prefix].shape[1] # 2048 device = ckpt_src[base_in_prefix].device dtype = ckpt_src[base_in_prefix].dtype exist_in_chans = [] # what can be loaded, per viewpoint exist_out_chans = [] # what can be loaded, per viewpoint want_in_chans = any4d_config.num_video_in_channels # what is desired, per viewpoint want_out_chans = any4d_config.num_video_out_channels # what is desired, per viewpoint for view_idx in range(V): in_prefix = in_prefixes[view_idx] out_prefix = out_prefixes[view_idx] cur_exist_in_chans = ckpt_src[in_prefix].shape[1] // PP if in_prefix in ckpt_src else 0 # 18 cur_exist_out_chans = ckpt_src[out_prefix].shape[0] // PP if out_prefix in ckpt_src else 0 # 16 cur_want_in_chans = want_in_chans[view_idx] # 18 cur_want_out_chans = want_out_chans[view_idx] # 18 exist_in_chans.append(cur_exist_in_chans) exist_out_chans.append(cur_exist_out_chans) log.info('') log.info(f'@@@@@@@@ for high-dim view {view_idx} proj: ' f'available: {cur_exist_in_chans} --> {cur_exist_out_chans} channels, ' f'requested: {cur_want_in_chans} --> {cur_want_out_chans} channels') ckpt_dst[in_prefix] = -0.1 * torch.ones( (in_embdim, cur_want_in_chans * PP), device=device, dtype=dtype) ckpt_dst[out_prefix] = -0.1 * torch.ones( (cur_want_out_chans * PP, out_embdim), device=device, dtype=dtype) # Now copy or zero-init or random-init or ignore each stream entry as needed for (key, (channel_start, channel_end, in_proj_strat, out_proj_strat)) in any4d_config.video_entries[view_idx].items(): friendly_entry = f'{key} (view {view_idx}, channels {channel_start}-{channel_end})' # Handle input projection weight logistics in_weights_available = (channel_end <= exist_in_chans[view_idx]) in_proj_init = in_proj_strat.split('/')[-1 if in_weights_available else 0] if in_proj_strat is not None else None # ^ e.g. 'zero/load' will become either 'zero' or 'load' if weights are new or exist respectively if in_proj_init == 'load' and not(in_weights_available): raise ValueError(f'ckpt does NOT have all weights available for {friendly_entry}, but in_proj_init = {in_proj_init}') if in_proj_init != 'load' and in_weights_available: log.warning(f'ckpt has all weights available for {friendly_entry}, but in_proj_init = {in_proj_init}, so these are NOT used') # NOTE(bvh): In PatchEmbed, there is rearrange b c (t r) (h m) (w n) -> b t h w (c r m n), # which means the ordering is channel major and patch index minor, so * PP is correct. if in_proj_init == 'load': log.info(f'loading in_proj for {friendly_entry} directly from ckpt') ckpt_dst[in_prefix][:, channel_start * PP:channel_end * PP] = \ ckpt_src[in_prefix][:, channel_start * PP:channel_end * PP] log.debug(f' dst slice: [:, {channel_start * PP}:{channel_end * PP}]') log.debug(f' src slice: [:, {channel_start * PP}:{channel_end * PP}]') elif in_proj_init == 'zero': log.info(f'zero-initializing in_proj for {friendly_entry}') ckpt_dst[in_prefix][:, channel_start * PP:channel_end * PP] = 0.0 # turn -0.1 into 0.0 log.debug(f' dst slice: [:, {channel_start * PP}:{channel_end * PP}]') elif in_proj_init == 'rand': log.info(f'random-initializing in_proj for {friendly_entry}') # Apply Kaiming uniform initialization but take total fan_in into account # NOTE(bvh): This randomness is NOT the same across GPUs / ranks. ckpt_dst[in_prefix][:, channel_start * PP:channel_end * PP] = \ kaiming_uniform_like(ckpt_dst[in_prefix][:, channel_start * PP:channel_end * PP], fan_in_total=cur_want_in_chans * PP) # torch.randn_like(ckpt_dst[in_prefix][:, channel_start * PP:channel_end * PP]) log.debug(f' dst slice: [:, {channel_start * PP}:{channel_end * PP}]') elif in_proj_init.startswith('copy:'): other_key = in_proj_init[5:] other_view = key2view[other_key] (other_c1, other_c2) = any4d_config.video_entries[other_view][other_key][0:2] other_friendly = f'{other_key} (view {other_view}, channels {other_c1}-{other_c2})' log.info(f'copying in_proj for {friendly_entry} from {other_friendly} in ckpt') other_prefix = in_prefixes[other_view] ckpt_dst[in_prefix][:, channel_start * PP:channel_end * PP] = \ ckpt_src[other_prefix][:, other_c1 * PP:other_c2 * PP] log.debug(f' dst slice: [:, {channel_start * PP}:{channel_end * PP}]') log.debug(f' src slice: [:, {other_c1 * PP}:{other_c2 * PP}]') else: raise ValueError(f'invalid in_proj_init for {friendly_entry}: {in_proj_init}') # Handle output projection weight logistics out_weights_available = (channel_end <= exist_out_chans[view_idx]) out_proj_init = out_proj_strat.split('/')[-1 if out_weights_available else 0] if out_proj_strat is not None else None # ^ e.g. 'zero/load' will become either 'zero' or 'load' if weights are new or exist respectively if out_proj_init == 'load' and not(out_weights_available): raise ValueError(f'ckpt does NOT have all weights available for {friendly_entry}, but out_proj_init = {out_proj_init}') if out_proj_init != 'load' and out_weights_available: log.warning(f'ckpt has all weights available for {friendly_entry}, but out_proj_init = {out_proj_init}, so these are NOT used') # NOTE(bvh): In unpatchify, there is rearrange B T H W (p1 p2 t C) -> B C (T t) (H p1) (W p2), # which is the INVERSE of PatchEmbed! This means the ordering is patch index major and channel minor, # so adding channels is more tricky (e.g. 0 1 x x 4 5 x x 8 9 x x ...) if out_proj_init is None: log.info(f'ignoring out_proj for {friendly_entry} (presumably unused; remains -0.1)') elif out_proj_init == 'load': log.info(f'loading out_proj for {friendly_entry} directly from ckpt') for p in range(PP): # ckpt_dst[out_prefix][channel_start * PP + p:channel_end * PP + p:PP, :] = \ # ckpt_src[out_prefix][channel_start * PP + p:channel_end * PP + p:PP, :] ckpt_dst[out_prefix][channel_start + p * cur_want_out_chans:channel_end + p * cur_want_out_chans, :] = \ ckpt_src[out_prefix][channel_start + p * cur_exist_out_chans:channel_end + p * cur_exist_out_chans, :] log.debug(f' dst slice: [{channel_start + p * cur_want_out_chans}:{channel_end + p * cur_want_out_chans}, :]') log.debug(f' src slice: [{channel_start + p * cur_exist_out_chans}:{channel_end + p * cur_exist_out_chans}, :]') elif out_proj_init == 'zero': log.info(f'zero-initializing out_proj for {friendly_entry}') for p in range(PP): # ckpt_dst[out_prefix][channel_start * PP + p:channel_end * PP + p:PP, :] = 0.0 # turn -0.1 into 0.0 ckpt_dst[out_prefix][channel_start + p * cur_want_out_chans:channel_end + p * cur_want_out_chans, :] = 0.0 # turn -0.1 into 0.0 log.debug(f' dst slice: [{channel_start + p * cur_want_out_chans}:{channel_end + p * cur_want_out_chans}, :]') elif out_proj_init == 'rand': log.info(f'random-initializing out_proj for {friendly_entry}') # Apply Kaiming uniform initialization but take total fan_in into account # NOTE(bvh): This randomness is NOT the same across GPUs / ranks. for p in range(PP): # ckpt_dst[out_prefix][channel_start * PP + p:channel_end * PP + p:PP, :] = \ # kaiming_uniform_like(ckpt_dst[out_prefix][channel_start:channel_end, :], # fan_in_total=out_embdim) # torch.randn_like(ckpt_dst[out_prefix][channel_start * PP + p:channel_end * PP + p:PP, :]) ckpt_dst[out_prefix][channel_start + p * cur_want_out_chans:channel_end + p * cur_want_out_chans, :] = \ kaiming_uniform_like(ckpt_dst[out_prefix][channel_start + p * cur_exist_out_chans:channel_end + p * cur_exist_out_chans, :], fan_in_total=out_embdim) log.debug(f' dst slice: [{channel_start + p * cur_want_out_chans}:{channel_end + p * cur_want_out_chans}, :]') elif out_proj_init.startswith('copy:'): other_key = out_proj_init[5:] other_view = key2view[other_key] (other_c1, other_c2) = any4d_config.video_entries[other_view][other_key][0:2] other_friendly = f'{other_key} (view {other_view}, channels {other_c1}-{other_c2})' log.info(f'copying out_proj for {friendly_entry} from {other_friendly} in ckpt') other_prefix = out_prefixes[other_view] for p in range(PP): # ckpt_dst[out_prefix][channel_start * PP + p:channel_end * PP + p:PP, :] = \ # ckpt_dst[other_prefix][other_c1 * PP + p:other_c2 * PP + p:PP, :] ckpt_dst[out_prefix][channel_start + p * cur_want_out_chans:channel_end + p * cur_want_out_chans, :] = \ ckpt_src[other_prefix][other_c1 + p * cur_exist_out_chans:other_c2 + p * cur_exist_out_chans, :] log.debug(f' dst slice: [{channel_start + p * cur_want_out_chans}:{channel_end + p * cur_want_out_chans}, :]') log.debug(f' src slice: [{other_c1 + p * cur_exist_out_chans}:{other_c2 + p * cur_exist_out_chans}, :]') else: raise ValueError(f'invalid out_proj_init for {friendly_entry}: {out_proj_init}') # Handle t_embedder_newviews -> always copy from t_embedder, or load if available for view_idx in range(1, V): if f't_embedder_newviews.{view_idx - 1}.1.linear_1.weight' in ckpt_src: log.info(f'loading t_embedder for view {view_idx} directly from ckpt') # do nothing (ckpt_dst is already a copy of ckpt_src) else: log.info(f'copying t_embedder for view {view_idx} from view 0') ckpt_dst[f't_embedder_newviews.{view_idx - 1}.1.linear_1.weight'] = \ ckpt_src['t_embedder.1.linear_1.weight'] ckpt_dst[f't_embedder_newviews.{view_idx - 1}.1.linear_2.weight'] = \ ckpt_src['t_embedder.1.linear_2.weight'] # ========================================================= # ======== Handle low-dimensional inputs & outputs ======== # ========================================================= # Handle self-attention weight logistics, based on lowdim_sattn_entries # TODO(bvh) # For now we do nothing which is also fine (default behavior = load when available, rand otherwise) # Handle cross-attention weight logistics, based on lowdim_xattn_entries # TODO(bvh) # For now we do nothing which is also fine (default behavior = load when available, rand otherwise) # Handle adaptive layer normalization weight logistics, based on lowdim_adaln_entries # We handle these manually: adaln_in1_fc1_weight = 'action_embedder_B_D.fc1.weight' adaln_in2_fc1_weight = 'action_embedder_B_3D.fc1.weight' # We do not touch these (so they are handled by pytorch): adaln_in1_fc1_bias = 'action_embedder_B_D.fc1.bias' adaln_in1_fc2_weight = 'action_embedder_B_D.fc2.weight' adaln_in1_fc2_bias = 'action_embedder_B_D.fc2.bias' adaln_in2_fc1_bias = 'action_embedder_B_3D.fc1.bias' adaln_in2_fc2_weight = 'action_embedder_B_3D.fc2.weight' adaln_in2_fc2_bias = 'action_embedder_B_3D.fc2.bias' if any4d_config.has_adaln_stream: # dit.action_embedder_B_D # Mlp( # (fc1): Linear(in_features=12, out_features=8192, bias=True) # (activation): GELU(approximate='tanh') # (fc2): Linear(in_features=8192, out_features=2048, bias=True) # (drop): Dropout(p=0, inplace=False) # ) # dit.action_embedder_B_3D # Mlp( # (fc1): Linear(in_features=12, out_features=8192, bias=True) # (activation): GELU(approximate='tanh') # (fc2): Linear(in_features=8192, out_features=6144, bias=True) # (drop): Dropout(p=0, inplace=False) # ) adaln_in1_embdim1 = dit.action_embedder_B_D.fc1.weight.shape[0] # 8192 adaln_in2_embdim1 = dit.action_embedder_B_3D.fc1.weight.shape[0] # 8192 # adaln_in1_embdim2 = dit.action_embedder_B_D.fc2.weight.shape[0] # 2048 # adaln_in2_embdim2 = dit.action_embedder_B_3D.fc2.weight.shape[0] # 6144 # adaln_exist_in_chans = ckpt_src[adaln_in1_fc1_weight].shape[1] if adaln_in1_fc1_weight in ckpt_src else 0 # adaln_want_in_chans = any4d_config.num_lowdim_adaln_channels adaln_exist_in_dim = ckpt_src[adaln_in1_fc1_weight].shape[1] if adaln_in1_fc1_weight in ckpt_src else 0 (Ta, Ca) = (any4d_config.num_lowdim_adaln_tokens, any4d_config.num_lowdim_adaln_channels) adaln_want_in_dim = Ta * Ca # NOTE(bvh): These embeddings are flattened! log.info('') log.info(f'@@@@@@@@ for low-dim adaln: ' f'available: {adaln_exist_in_dim} dim, ' f'requested: {adaln_want_in_dim} dim') ckpt_dst[adaln_in1_fc1_weight] = -0.1 * torch.ones( (adaln_in1_embdim1, adaln_want_in_dim), device=device, dtype=dtype) ckpt_dst[adaln_in2_fc1_weight] = -0.1 * torch.ones( (adaln_in2_embdim1, adaln_want_in_dim), device=device, dtype=dtype) for (key, (seq_start, seq_end, channel_start, channel_end, in_proj_strat)) in any4d_config.lowdim_adaln_entries.items(): friendly_entry = f'adaln {key} (tokens {seq_start}-{seq_end}, channels {channel_start}-{channel_end})' # Handle input projection weight logistics # in_weights_available = (seq_end * channel_end <= adaln_exist_in_dim) # wrong until 05/19/2026 # update 05/19: unlike videos, always consider/apply ALL channels in_weights_available = (seq_end * Ca <= adaln_exist_in_dim) # NOTE; ^ this may mean we are performing some duplicate operations below # (for overlapping tokens) but that's fine. in_proj_init = in_proj_strat.split('/')[-1 if in_weights_available else 0] if in_proj_strat is not None else None # ^ e.g. 'zero/load' will become either 'zero' or 'load' if weights are new or exist respectively if in_proj_init == 'load' and not(in_weights_available): # raise ValueError(f'ckpt does NOT have all weights available for {friendly_entry}, but in_proj_init = {in_proj_init}') # update 05/19: increase tolerance to vary world model horizons more gracefully log.warning(f'ckpt does NOT have all weights available for {friendly_entry}, but in_proj_init = {in_proj_init}, ' f'so we load as many tokens as available') if in_proj_init != 'load' and in_weights_available: log.warning(f'ckpt has all weights available for {friendly_entry}, but in_proj_init = {in_proj_init}, ' f'so these are NOT used') if in_proj_init == 'load': # WARNING(bvh): We have to assume that ckpt_src and ckpt_dst have the same # channels! # (however, different # tokens are allowed) log.info(f'loading in_proj for {friendly_entry} directly from ckpt (selected tokens only, but ALL channels)') lo = min(seq_start * Ca, adaln_exist_in_dim) hi = min(seq_end * Ca, adaln_exist_in_dim) log.info(f'lo={lo}, hi={hi}') ckpt_dst[adaln_in1_fc1_weight][:, lo:hi] = \ ckpt_src[adaln_in1_fc1_weight][:, lo:hi] ckpt_dst[adaln_in2_fc1_weight][:, lo:hi] = \ ckpt_src[adaln_in2_fc1_weight][:, lo:hi] # Leave other weights intact (= either random or load from ckpt_src) elif in_proj_init == 'zero': log.info(f'zero-initializing in_proj for {friendly_entry} (selected tokens only, but ALL channels)') ckpt_dst[adaln_in1_fc1_weight][:, seq_start * Ca:seq_end * Ca] = 0.0 ckpt_dst[adaln_in2_fc1_weight][:, seq_start * Ca:seq_end * Ca] = 0.0 # Leave other weights intact (= either random or load from ckpt_src) elif in_proj_init == 'rand': log.info(f'random-initializing in_proj for {friendly_entry} (selected tokens only, but ALL channels)') # Apply Kaiming uniform initialization but take total fan_in into account # NOTE(bvh): This randomness is NOT the same across GPUs / ranks. ckpt_dst[adaln_in1_fc1_weight][:, seq_start * Ca:seq_end * Ca] = \ kaiming_uniform_like(ckpt_dst[adaln_in1_fc1_weight][:, seq_start * Ca:seq_end * Ca], fan_in_total=adaln_want_in_dim) ckpt_dst[adaln_in2_fc1_weight][:, seq_start * Ca:seq_end * Ca] = \ kaiming_uniform_like(ckpt_dst[adaln_in2_fc1_weight][:, seq_start * Ca:seq_end * Ca], fan_in_total=adaln_want_in_dim) elif in_proj_init.startswith('copy:'): other_key = in_proj_init[5:] other_view = key2view[other_key] (other_s1, other_s2, other_c1, other_c2) = any4d_config.lowdim_adaln_entries[other_key][0:4] other_friendly = f'adaln {other_key} (tokens {other_s1}-{other_s2}, channels {other_c1}-{other_c2})' log.info(f'copying in_proj for {friendly_entry} from {other_friendly} in ckpt (selected tokens only, but ALL channels)') other_prefix = adaln_in1_fc1_weight ckpt_dst[adaln_in1_fc1_weight][:, seq_start * Ca:seq_end * Ca] = \ ckpt_src[other_prefix][:, other_s1 * Ca:other_s2 * Ca] other_prefix = adaln_in2_fc1_weight ckpt_dst[adaln_in2_fc1_weight][:, seq_start * Ca:seq_end * Ca] = \ ckpt_src[other_prefix][:, other_s1 * Ca:other_s2 * Ca] # Leave other weights intact (= either random or load from ckpt_src) # Finally, update the entire checkpoint (including transformer blocks) with the new weights before loading them ckpt.update(ckpt_dst) # NOTE(bvh): There should be no meta tensors before or after this point. # import random; a = random.choice(list(dit.parameters())) (missing_keys, surplus_keys) = dit.load_state_dict(ckpt, strict=False, assign=True) log.warning(f'all missing keys: {missing_keys}') log.warning(f'all surplus keys: {surplus_keys}') # Typically something like: ['action_embedder_B_D.fc1.bias', 'action_embedder_B_D.fc2.weight', # 'action_embedder_B_D.fc2.bias', 'action_embedder_B_3D.fc1.bias', 'action_embedder_B_3D.fc2.weight', # 'action_embedder_B_3D.fc2.bias', 'view_embs_newviews.0', # 'final_layer_newviews.0.adaln_modulation.1.weight', 'final_layer_newviews.0.adaln_modulation.2.weight']. # Checking this is important to ensure no oversights in surgery. expected_missing_keys = [x for x in missing_keys if x in [ adaln_in1_fc1_bias, adaln_in1_fc2_weight, adaln_in1_fc2_bias, adaln_in2_fc1_bias, adaln_in2_fc2_weight, adaln_in2_fc2_bias] or 'view_embs_newviews' in x # or 't_embedder_newviews' in x # we always copy manually hence should never be missing or ('final_layer_newviews' in x and 'adaln_modulation' in x)] unexpected_missing_keys = [x for x in missing_keys if x not in expected_missing_keys] # sattn_rope is the removed dead param (see a4d_network.py); old ckpts still carry it -> surplus. expected_surplus_keys = [x for x in surplus_keys if 'cross_attn.attn_op._extra_state' in x or 'sattn_rope' in x] unexpected_surplus_keys = [x for x in surplus_keys if x not in expected_surplus_keys] if len(unexpected_missing_keys) > 0: log.error(f'unexpected_missing_keys: {unexpected_missing_keys}') else: log.success(f'no unexpected missing keys (i.e. those intentionally left to default initialization) found') if len(unexpected_surplus_keys) > 0: log.error(f'unexpected_surplus_keys: {unexpected_surplus_keys}') else: log.success(f'no unexpected surplus keys (i.e. without cross_attn.attn_op._extra_state) found') # assert len(missing_keys) == 0, f'missing_keys: {missing_keys}' # TODO(bvh): Given divergent random inits across ranks, # verify whether they synchronize automatically at some point? # ^ answer: yes return dit ''' dit.view_embs_newviews[0].v Parameter[1, 2048] 8Kb x∈[-3.234, 3.457] μ=0.013 σ=1.026 grad tensor([[-0.701, 0.310, -1.422, ..., -0.149, 0.449, -0.236]], requires_grad=True) Parameter[1, 2048] 8Kb x∈[-3.351, 3.301] μ=0.015 σ=0.994 grad tensor([[-0.568, 0.968, 0.356, ..., -0.988, 1.117, -0.092]], requires_grad=True) '''