egolm-protocol-v2-code / mGPT /utils /load_checkpoint.py

EgoLM baseline code (Ego3DLM snapshot, unmodified) + upload notes

3de4238 verified about 1 month ago

7.16 kB

import torch
import torch
from collections import OrderedDict

def load_t5_from_3dllm_ckpt_V2(ckpt_path, t5, device="cpu"):
    obj = torch.load(ckpt_path, map_location=device, weights_only=False)
    sd = obj.get("model", obj)

    def fit_vocab(src, tgt_shape):
        if list(src.shape) == list(tgt_shape):
            return src
        Vsrc, D = src.shape
        Vtgt, D2 = tgt_shape
        assert D == D2, f"embed dim mismatch {D} vs {D2}"
        if Vsrc > Vtgt:
            raise ValueError(f"Checkpoint vocab {Vsrc} > model vocab {Vtgt}")
        pad = torch.zeros((Vtgt - Vsrc, D), dtype=src.dtype, device=src.device)
        return torch.cat([src, pad], dim=0)

    new_sd = OrderedDict()
    for name, param in t5.state_dict().items():
        ckpt_name = name if name in sd else f"t5_model.{name}"
        if ckpt_name not in sd:
            continue
        w = sd[ckpt_name]
        if "embed_tokens.weight" in name or "shared.weight" in name or "lm_head.weight" in name:
            w = fit_vocab(w, param.shape)
        new_sd[name] = w

    t5.load_state_dict(new_sd, strict=False)
    t5.tie_weights()
    return t5


def load_t5_from_3dllm_ckpt(ckpt_path, t5, device="cpu"):
    
    obj = torch.load(ckpt_path, map_location=device, weights_only=False)
    sd = obj.get("model", obj)

    def take(key):
        # 3d llm keys start with : 't5_model.xxx'
        k = key if key in sd else f"t5_model.{key}"
        return sd[k]
    def fit_vocab(src, tgt_shape):  # (Vsrc, D) -> (Vtgt, D)
        if list(src.shape) == list(tgt_shape):
            return src
        Vsrc, D = src.shape
        Vtgt, D2 = tgt_shape
        assert D == D2, f"embed dim mismatch {D} vs {D2}"
        if Vsrc >= Vtgt:  #
            raise Exception('the source(checkpoint) must have smaller vocab size')
        pad = torch.zeros((Vtgt - Vsrc, D), dtype=src.dtype)
        return torch.cat([src, pad], dim=0)

    with torch.no_grad():
        # shared / embed_tokens / lm_head
        sw = take("shared.weight")
        t5.shared.weight.copy_(fit_vocab(sw, t5.shared.weight.shape))

        ew = take("encoder.embed_tokens.weight")
        t5.encoder.embed_tokens.weight.copy_(fit_vocab(ew, t5.encoder.embed_tokens.weight.shape))

        dw = take("decoder.embed_tokens.weight")
        t5.decoder.embed_tokens.weight.copy_(fit_vocab(dw, t5.decoder.embed_tokens.weight.shape))

        lw = take("lm_head.weight")
        t5.lm_head.weight.copy_(fit_vocab(lw, t5.lm_head.weight.shape))

    t5.tie_weights()
    return t5
    
def load_pretrained(cfg, model, logger=None, phase="train"):
    if logger is not None:
        logger.info(f"Loading pretrain model from {cfg.TRAIN.PRETRAINED}")

    if phase == "train":
        ckpt_path = cfg.TRAIN.PRETRAINED
    elif phase == "test":
        ckpt_path = cfg.TEST.CHECKPOINTS

    state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False)["state_dict"]

    ignore_prefixes = []

    # dangerous modification
    if cfg.TRAIN.get('LOAD_CKPT_WO_TP_ENC', False):
        ignore_prefixes.append("three_points_proj.")

    # GRPO validation: exclude ref_model weights from checkpoint
    if cfg.TRAIN.get('GRPO_VAL', False):
        ignore_prefixes.append("ref_model.")

    if ignore_prefixes:
        state_dict = {k: v for k, v in state_dict.items()
                      if not any(k.startswith(p) for p in ignore_prefixes)}
        if logger is not None:
            logger.info(f"Filtered out keys with prefixes: {ignore_prefixes}")
        model.load_state_dict(state_dict, strict=False)
    else:
        model.load_state_dict(state_dict, strict=False)
    return model


def load_pretrained_vae(cfg, model, logger=None):
    state_dict = torch.load(cfg.TRAIN.PRETRAINED_VAE,
                            map_location="cpu")['state_dict']
    if logger is not None:
        logger.info(f"Loading pretrain vae from {cfg.TRAIN.PRETRAINED_VAE}")
        
    # Extract encoder/decoder
    from collections import OrderedDict
    vae_dict = OrderedDict()
    for k, v in state_dict.items():
        if "motion_vae" in k:
            name = k.replace("motion_vae.", "")
            vae_dict[name] = v
        elif "vae" in k:
            name = k.replace("vae.", "")
            vae_dict[name] = v
    if hasattr(model, 'vae'):
        model.vae.load_state_dict(vae_dict, strict=True)
    else:
        model.motion_vae.load_state_dict(vae_dict, strict=True)
    
    return model


def load_pretrained_grpo_ref(cfg, model, logger=None):
    """Load GRPO reference model from checkpoint.
    Reference model is co-located with the main model on each rank's device.
    The legacy ``GRPO_REF_DEVICE`` setting is intentionally ignored: with DDP,
    Lightning recursively moves child submodules to the per-rank device, so any
    fixed device string would go stale and trigger cross-device tensor moves.
    """
    ref_path = cfg.TRAIN.GRPO_REF_MODEL_PATH
    ref_device = model.device
    if logger is not None:
        logger.info(f"Loading GRPO reference model from {ref_path}")
        logger.info(f"Reference model will be placed on {ref_device} (co-located with main model)")

    from mGPT.models.build_model import build_model

    # Create ref model with same architecture as training model
    ref_model = build_model(cfg, model.datamodule)

    # Add same special tokens as training model
    if cfg.model.params.get("dummy_motion", False):
        ref_model.lm.init_dummy_motion_token(
            cfg.model.params.p_dummy,
            cfg.model.params.per_token,
            cfg.model.params.get('use_single_span', False),
            cfg.model.params.get('span_pairs', 6),
            cfg.model.params.get('p_span', 1.0)
        )
    if cfg.model.params.get("tp_special_token", False):
        ref_model.lm.init_tp_special_token()
    if cfg.LOSS.ABLATION.get("TRAIN_ALL_AT_ONCE", False):
        ref_model.lm.init_text_special_token()
    if cfg.TRAIN.get("CHAIN_OF_THOUGHT", False):
        ref_model.lm.init_cot_special_token()

    # Load state dict (vocab size now matches checkpoint)
    state_dict = torch.load(ref_path, map_location="cpu", weights_only=False)["state_dict"]
    # If the ckpt itself was a GRPO save (contains its own nested ref_model.*
    # subtree), strip that prefix — those keys would otherwise be unexpected
    # in a fresh MotionGPT (= ref_model) which has no nested ref_model. Same
    # for three_points_proj when LOAD_CKPT_WO_TP_ENC is set.
    ignore_prefixes = ["ref_model."]
    if cfg.TRAIN.get('LOAD_CKPT_WO_TP_ENC', False):
        ignore_prefixes.append("three_points_proj.")
    state_dict = {k: v for k, v in state_dict.items()
                  if not any(k.startswith(p) for p in ignore_prefixes)}
    # strict=False to also tolerate any benign add/remove (e.g. peft wrapping
    # differences when the ckpt was full-FT but config wraps with LoRA).
    ref_model.load_state_dict(state_dict, strict=False)

    ref_model = ref_model.to(ref_device)
    ref_model.eval()

    # Freeze all parameters
    for param in ref_model.parameters():
        param.requires_grad = False

    model.ref_model = ref_model

    return model

	import torch
	import torch
	from collections import OrderedDict

	def load_t5_from_3dllm_ckpt_V2(ckpt_path, t5, device="cpu"):
	obj = torch.load(ckpt_path, map_location=device, weights_only=False)
	sd = obj.get("model", obj)

	def fit_vocab(src, tgt_shape):
	if list(src.shape) == list(tgt_shape):
	return src
	Vsrc, D = src.shape
	Vtgt, D2 = tgt_shape
	assert D == D2, f"embed dim mismatch {D} vs {D2}"
	if Vsrc > Vtgt:
	raise ValueError(f"Checkpoint vocab {Vsrc} > model vocab {Vtgt}")
	pad = torch.zeros((Vtgt - Vsrc, D), dtype=src.dtype, device=src.device)
	return torch.cat([src, pad], dim=0)

	new_sd = OrderedDict()
	for name, param in t5.state_dict().items():
	ckpt_name = name if name in sd else f"t5_model.{name}"
	if ckpt_name not in sd:
	continue
	w = sd[ckpt_name]
	if "embed_tokens.weight" in name or "shared.weight" in name or "lm_head.weight" in name:
	w = fit_vocab(w, param.shape)
	new_sd[name] = w

	t5.load_state_dict(new_sd, strict=False)
	t5.tie_weights()
	return t5


	def load_t5_from_3dllm_ckpt(ckpt_path, t5, device="cpu"):

	obj = torch.load(ckpt_path, map_location=device, weights_only=False)
	sd = obj.get("model", obj)

	def take(key):
	# 3d llm keys start with : 't5_model.xxx'
	k = key if key in sd else f"t5_model.{key}"
	return sd[k]
	def fit_vocab(src, tgt_shape): # (Vsrc, D) -> (Vtgt, D)
	if list(src.shape) == list(tgt_shape):
	return src
	Vsrc, D = src.shape
	Vtgt, D2 = tgt_shape
	assert D == D2, f"embed dim mismatch {D} vs {D2}"
	if Vsrc >= Vtgt: #
	raise Exception('the source(checkpoint) must have smaller vocab size')
	pad = torch.zeros((Vtgt - Vsrc, D), dtype=src.dtype)
	return torch.cat([src, pad], dim=0)

	with torch.no_grad():
	# shared / embed_tokens / lm_head
	sw = take("shared.weight")
	t5.shared.weight.copy_(fit_vocab(sw, t5.shared.weight.shape))

	ew = take("encoder.embed_tokens.weight")
	t5.encoder.embed_tokens.weight.copy_(fit_vocab(ew, t5.encoder.embed_tokens.weight.shape))

	dw = take("decoder.embed_tokens.weight")
	t5.decoder.embed_tokens.weight.copy_(fit_vocab(dw, t5.decoder.embed_tokens.weight.shape))

	lw = take("lm_head.weight")
	t5.lm_head.weight.copy_(fit_vocab(lw, t5.lm_head.weight.shape))

	t5.tie_weights()
	return t5

	def load_pretrained(cfg, model, logger=None, phase="train"):
	if logger is not None:
	logger.info(f"Loading pretrain model from {cfg.TRAIN.PRETRAINED}")

	if phase == "train":
	ckpt_path = cfg.TRAIN.PRETRAINED
	elif phase == "test":
	ckpt_path = cfg.TEST.CHECKPOINTS

	state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False)["state_dict"]

	ignore_prefixes = []

	# dangerous modification
	if cfg.TRAIN.get('LOAD_CKPT_WO_TP_ENC', False):
	ignore_prefixes.append("three_points_proj.")

	# GRPO validation: exclude ref_model weights from checkpoint
	if cfg.TRAIN.get('GRPO_VAL', False):
	ignore_prefixes.append("ref_model.")

	if ignore_prefixes:
	state_dict = {k: v for k, v in state_dict.items()
	if not any(k.startswith(p) for p in ignore_prefixes)}
	if logger is not None:
	logger.info(f"Filtered out keys with prefixes: {ignore_prefixes}")
	model.load_state_dict(state_dict, strict=False)
	else:
	model.load_state_dict(state_dict, strict=False)
	return model


	def load_pretrained_vae(cfg, model, logger=None):
	state_dict = torch.load(cfg.TRAIN.PRETRAINED_VAE,
	map_location="cpu")['state_dict']
	if logger is not None:
	logger.info(f"Loading pretrain vae from {cfg.TRAIN.PRETRAINED_VAE}")

	# Extract encoder/decoder
	from collections import OrderedDict
	vae_dict = OrderedDict()
	for k, v in state_dict.items():
	if "motion_vae" in k:
	name = k.replace("motion_vae.", "")
	vae_dict[name] = v
	elif "vae" in k:
	name = k.replace("vae.", "")
	vae_dict[name] = v
	if hasattr(model, 'vae'):
	model.vae.load_state_dict(vae_dict, strict=True)
	else:
	model.motion_vae.load_state_dict(vae_dict, strict=True)

	return model


	def load_pretrained_grpo_ref(cfg, model, logger=None):
	"""Load GRPO reference model from checkpoint.

	Reference model is co-located with the main model on each rank's device.
	The legacy ``GRPO_REF_DEVICE`` setting is intentionally ignored: with DDP,
	Lightning recursively moves child submodules to the per-rank device, so any
	fixed device string would go stale and trigger cross-device tensor moves.
	"""
	ref_path = cfg.TRAIN.GRPO_REF_MODEL_PATH
	ref_device = model.device
	if logger is not None:
	logger.info(f"Loading GRPO reference model from {ref_path}")
	logger.info(f"Reference model will be placed on {ref_device} (co-located with main model)")

	from mGPT.models.build_model import build_model

	# Create ref model with same architecture as training model
	ref_model = build_model(cfg, model.datamodule)

	# Add same special tokens as training model
	if cfg.model.params.get("dummy_motion", False):
	ref_model.lm.init_dummy_motion_token(
	cfg.model.params.p_dummy,
	cfg.model.params.per_token,
	cfg.model.params.get('use_single_span', False),
	cfg.model.params.get('span_pairs', 6),
	cfg.model.params.get('p_span', 1.0)
	)
	if cfg.model.params.get("tp_special_token", False):
	ref_model.lm.init_tp_special_token()
	if cfg.LOSS.ABLATION.get("TRAIN_ALL_AT_ONCE", False):
	ref_model.lm.init_text_special_token()
	if cfg.TRAIN.get("CHAIN_OF_THOUGHT", False):
	ref_model.lm.init_cot_special_token()

	# Load state dict (vocab size now matches checkpoint)
	state_dict = torch.load(ref_path, map_location="cpu", weights_only=False)["state_dict"]
	# If the ckpt itself was a GRPO save (contains its own nested ref_model.*
	# subtree), strip that prefix — those keys would otherwise be unexpected
	# in a fresh MotionGPT (= ref_model) which has no nested ref_model. Same
	# for three_points_proj when LOAD_CKPT_WO_TP_ENC is set.
	ignore_prefixes = ["ref_model."]
	if cfg.TRAIN.get('LOAD_CKPT_WO_TP_ENC', False):
	ignore_prefixes.append("three_points_proj.")
	state_dict = {k: v for k, v in state_dict.items()
	if not any(k.startswith(p) for p in ignore_prefixes)}
	# strict=False to also tolerate any benign add/remove (e.g. peft wrapping
	# differences when the ckpt was full-FT but config wraps with LoRA).
	ref_model.load_state_dict(state_dict, strict=False)

	ref_model = ref_model.to(ref_device)
	ref_model.eval()

	# Freeze all parameters
	for param in ref_model.parameters():
	param.requires_grad = False

	model.ref_model = ref_model

	return model