MiniMax-H3-Semantic-Bridge / research /raw_scripts /train_embedding_bridge_test.py

Initial release of MiniMax H3 Semantic Bridge v1.0

752626b verified about 2 months ago

8.85 kB

import os
import math
import torch

from safetensors import safe_open
from safetensors.torch import load_file, save_file


# =========================================================
# PATHS
# =========================================================

H3_EMB = r"E:\Models SD XL\UNET\Minimax_H3\h3_qwen3vl32b_embeddings_bf16.safetensors"

SENSENOVA = r"D:\ComfyUI_Python312\ComfyUI\models\unet\SenseNova-U1.5-8B-MoT-pruned-bf16.safetensors"

OUT = r"E:\Models SD XL\UNET\Minimax_H3\sensenova_to_h3_embedding_bridge_rank512.safetensors"


# =========================================================
# SETTINGS
# =========================================================

INPUT_DIM = 4096
OUTPUT_DIM = 5120
RANK = 512

BATCH_SIZE = 256
EPOCHS = 5
LR = 1e-3

VAL_SIZE = 12000

SEED = 1234


# =========================================================
# DEVICE
# =========================================================

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("=" * 80)
print("SenseNova -> MiniMax H3 embedding bridge test")
print("=" * 80)

print("Device:", device)

if device.type == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))
    print(
        "VRAM:",
        round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2),
        "GB"
    )


# =========================================================
# LOAD H3 EMBEDDINGS
# =========================================================

print()
print("Loading H3 embeddings...")

h3_file = load_file(H3_EMB, device="cpu")

h3 = h3_file["model.embed_tokens.weight"]

print("H3 shape:", tuple(h3.shape))
print("H3 dtype:", h3.dtype)


# =========================================================
# LOAD SENSENOVA EMBEDDINGS ONLY
# =========================================================

print()
print("Loading SenseNova token embeddings...")

with safe_open(SENSENOVA, framework="pt", device="cpu") as f:

    key = "language_model.model.embed_tokens.weight"

    if key not in f.keys():
        raise RuntimeError(
            f"SenseNova embedding tensor not found: {key}"
        )

    sense = f.get_tensor(key)


print("SenseNova shape:", tuple(sense.shape))
print("SenseNova dtype:", sense.dtype)


# =========================================================
# VALIDATE
# =========================================================

if h3.shape[0] != sense.shape[0]:
    raise RuntimeError(
        f"Vocabulary mismatch: H3={h3.shape[0]}, SenseNova={sense.shape[0]}"
    )

if sense.shape[1] != INPUT_DIM:
    raise RuntimeError(
        f"Unexpected SenseNova hidden size: {sense.shape[1]}"
    )

if h3.shape[1] != OUTPUT_DIM:
    raise RuntimeError(
        f"Unexpected H3 hidden size: {h3.shape[1]}"
    )


VOCAB = h3.shape[0]

print()
print("Vocabulary rows:", VOCAB)


# =========================================================
# NORMALIZATION
# =========================================================

# Convert storage to float32 only per batch later.
# Calculate global average norms using manageable chunks.

def average_norm(x, chunk=4096):

    total = 0.0
    count = 0

    for start in range(0, x.shape[0], chunk):

        end = min(start + chunk, x.shape[0])

        b = x[start:end].float()

        total += b.norm(dim=1).sum().item()
        count += b.shape[0]

    return total / count


print()
print("Calculating embedding norm statistics...")

sense_norm = average_norm(sense)
h3_norm = average_norm(h3)

print("SenseNova avg norm:", sense_norm)
print("H3 avg norm:       ", h3_norm)


# =========================================================
# TRAIN / VALIDATION SPLIT
# =========================================================

g = torch.Generator()
g.manual_seed(SEED)

perm = torch.randperm(VOCAB, generator=g)

val_ids = perm[:VAL_SIZE]
train_ids = perm[VAL_SIZE:]

print()
print("Train rows:", len(train_ids))
print("Validation rows:", len(val_ids))


# =========================================================
# LOW-RANK PROJECTOR
#
# 4096 -> 512 -> 5120
# no activation: this is a low-rank linear mapping
# =========================================================

A = torch.nn.Linear(
    INPUT_DIM,
    RANK,
    bias=False,
    device=device,
    dtype=torch.float32,
)

B = torch.nn.Linear(
    RANK,
    OUTPUT_DIM,
    bias=False,
    device=device,
    dtype=torch.float32,
)


torch.nn.init.normal_(
    A.weight,
    std=1.0 / math.sqrt(INPUT_DIM)
)

torch.nn.init.zeros_(B.weight)


params = list(A.parameters()) + list(B.parameters())

optimizer = torch.optim.AdamW(
    params,
    lr=LR,
    weight_decay=1e-4,
)


# =========================================================
# LOSS
# =========================================================

def bridge_forward(x):

    return B(A(x))


def loss_fn(pred, target):

    # Directional similarity.
    cosine = 1.0 - torch.nn.functional.cosine_similarity(
        pred,
        target,
        dim=-1
    ).mean()

    # Normalize magnitude sensitivity.
    pred_n = torch.nn.functional.normalize(pred, dim=-1)
    target_n = torch.nn.functional.normalize(target, dim=-1)

    mse = torch.nn.functional.mse_loss(
        pred_n,
        target_n
    )

    return cosine + mse


# =========================================================
# EVALUATION
# =========================================================

@torch.no_grad()
def evaluate(ids):

    A.eval()
    B.eval()

    cos_sum = 0.0
    mse_sum = 0.0
    count = 0

    for start in range(0, len(ids), BATCH_SIZE):

        idx = ids[start:start + BATCH_SIZE]

        x = sense[idx].float().to(device)
        y = h3[idx].float().to(device)

        pred = bridge_forward(x)

        cos = torch.nn.functional.cosine_similarity(
            pred,
            y,
            dim=-1
        )

        mse = torch.nn.functional.mse_loss(
            pred,
            y,
            reduction="none"
        ).mean(dim=1)

        cos_sum += cos.sum().item()
        mse_sum += mse.sum().item()

        count += len(idx)

    return cos_sum / count, mse_sum / count


# =========================================================
# BASELINE
# =========================================================

print()
print("=" * 80)
print("TRAINING")
print("=" * 80)


# =========================================================
# TRAIN
# =========================================================

for epoch in range(EPOCHS):

    A.train()
    B.train()

    shuffled = train_ids[
        torch.randperm(len(train_ids), generator=g)
    ]

    running = 0.0
    batches = 0

    for start in range(0, len(shuffled), BATCH_SIZE):

        idx = shuffled[start:start + BATCH_SIZE]

        x = sense[idx].float().to(device)
        y = h3[idx].float().to(device)

        optimizer.zero_grad(set_to_none=True)

        pred = bridge_forward(x)

        loss = loss_fn(pred, y)

        loss.backward()

        torch.nn.utils.clip_grad_norm_(
            params,
            1.0
        )

        optimizer.step()

        running += loss.item()
        batches += 1

        if batches % 100 == 0:

            print(
                f"Epoch {epoch + 1}/{EPOCHS} | "
                f"batch {batches} | "
                f"loss {running / batches:.6f}"
            )

    val_cos, val_mse = evaluate(val_ids)

    print()
    print(
        f"EPOCH {epoch + 1} COMPLETE | "
        f"train_loss={running / batches:.6f} | "
        f"val_cos={val_cos:.6f} | "
        f"val_mse={val_mse:.8f}"
    )
    print()


# =========================================================
# FINAL EVALUATION
# =========================================================

val_cos, val_mse = evaluate(val_ids)

print("=" * 80)
print("FINAL")
print("=" * 80)

print("Validation cosine similarity:", val_cos)
print("Validation MSE:", val_mse)


# =========================================================
# SAVE
# =========================================================

save_file(
    {
        "proj_in.weight": A.weight.detach().cpu().to(torch.float16),
        "proj_out.weight": B.weight.detach().cpu().to(torch.float16),
    },
    OUT,
    metadata={
        "source": "SenseNova-U1.5-8B-MoT",
        "target": "MiniMax-H3-Qwen3VL32B",
        "input_dim": str(INPUT_DIM),
        "rank": str(RANK),
        "output_dim": str(OUTPUT_DIM),
        "validation_cosine": str(val_cos),
        "validation_mse": str(val_mse),
    }
)

print()
print("Saved bridge:")
print(OUT)
print()
print("DONE")

	import os
	import math
	import torch

	from safetensors import safe_open
	from safetensors.torch import load_file, save_file


	# =========================================================
	# PATHS
	# =========================================================

	H3_EMB = r"E:\Models SD XL\UNET\Minimax_H3\h3_qwen3vl32b_embeddings_bf16.safetensors"

	SENSENOVA = r"D:\ComfyUI_Python312\ComfyUI\models\unet\SenseNova-U1.5-8B-MoT-pruned-bf16.safetensors"

	OUT = r"E:\Models SD XL\UNET\Minimax_H3\sensenova_to_h3_embedding_bridge_rank512.safetensors"


	# =========================================================
	# SETTINGS
	# =========================================================

	INPUT_DIM = 4096
	OUTPUT_DIM = 5120
	RANK = 512

	BATCH_SIZE = 256
	EPOCHS = 5
	LR = 1e-3

	VAL_SIZE = 12000

	SEED = 1234


	# =========================================================
	# DEVICE
	# =========================================================

	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	print("=" * 80)
	print("SenseNova -> MiniMax H3 embedding bridge test")
	print("=" * 80)

	print("Device:", device)

	if device.type == "cuda":
	print("GPU:", torch.cuda.get_device_name(0))
	print(
	"VRAM:",
	round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2),
	"GB"
	)


	# =========================================================
	# LOAD H3 EMBEDDINGS
	# =========================================================

	print()
	print("Loading H3 embeddings...")

	h3_file = load_file(H3_EMB, device="cpu")

	h3 = h3_file["model.embed_tokens.weight"]

	print("H3 shape:", tuple(h3.shape))
	print("H3 dtype:", h3.dtype)


	# =========================================================
	# LOAD SENSENOVA EMBEDDINGS ONLY
	# =========================================================

	print()
	print("Loading SenseNova token embeddings...")

	with safe_open(SENSENOVA, framework="pt", device="cpu") as f:

	key = "language_model.model.embed_tokens.weight"

	if key not in f.keys():
	raise RuntimeError(
	f"SenseNova embedding tensor not found: {key}"
	)

	sense = f.get_tensor(key)


	print("SenseNova shape:", tuple(sense.shape))
	print("SenseNova dtype:", sense.dtype)


	# =========================================================
	# VALIDATE
	# =========================================================

	if h3.shape[0] != sense.shape[0]:
	raise RuntimeError(
	f"Vocabulary mismatch: H3={h3.shape[0]}, SenseNova={sense.shape[0]}"
	)

	if sense.shape[1] != INPUT_DIM:
	raise RuntimeError(
	f"Unexpected SenseNova hidden size: {sense.shape[1]}"
	)

	if h3.shape[1] != OUTPUT_DIM:
	raise RuntimeError(
	f"Unexpected H3 hidden size: {h3.shape[1]}"
	)


	VOCAB = h3.shape[0]

	print()
	print("Vocabulary rows:", VOCAB)


	# =========================================================
	# NORMALIZATION
	# =========================================================

	# Convert storage to float32 only per batch later.
	# Calculate global average norms using manageable chunks.

	def average_norm(x, chunk=4096):

	total = 0.0
	count = 0

	for start in range(0, x.shape[0], chunk):

	end = min(start + chunk, x.shape[0])

	b = x[start:end].float()

	total += b.norm(dim=1).sum().item()
	count += b.shape[0]

	return total / count


	print()
	print("Calculating embedding norm statistics...")

	sense_norm = average_norm(sense)
	h3_norm = average_norm(h3)

	print("SenseNova avg norm:", sense_norm)
	print("H3 avg norm: ", h3_norm)


	# =========================================================
	# TRAIN / VALIDATION SPLIT
	# =========================================================

	g = torch.Generator()
	g.manual_seed(SEED)

	perm = torch.randperm(VOCAB, generator=g)

	val_ids = perm[:VAL_SIZE]
	train_ids = perm[VAL_SIZE:]

	print()
	print("Train rows:", len(train_ids))
	print("Validation rows:", len(val_ids))


	# =========================================================
	# LOW-RANK PROJECTOR
	#
	# 4096 -> 512 -> 5120
	# no activation: this is a low-rank linear mapping
	# =========================================================

	A = torch.nn.Linear(
	INPUT_DIM,
	RANK,
	bias=False,
	device=device,
	dtype=torch.float32,
	)

	B = torch.nn.Linear(
	RANK,
	OUTPUT_DIM,
	bias=False,
	device=device,
	dtype=torch.float32,
	)


	torch.nn.init.normal_(
	A.weight,
	std=1.0 / math.sqrt(INPUT_DIM)
	)

	torch.nn.init.zeros_(B.weight)


	params = list(A.parameters()) + list(B.parameters())

	optimizer = torch.optim.AdamW(
	params,
	lr=LR,
	weight_decay=1e-4,
	)


	# =========================================================
	# LOSS
	# =========================================================

	def bridge_forward(x):

	return B(A(x))


	def loss_fn(pred, target):

	# Directional similarity.
	cosine = 1.0 - torch.nn.functional.cosine_similarity(
	pred,
	target,
	dim=-1
	).mean()

	# Normalize magnitude sensitivity.
	pred_n = torch.nn.functional.normalize(pred, dim=-1)
	target_n = torch.nn.functional.normalize(target, dim=-1)

	mse = torch.nn.functional.mse_loss(
	pred_n,
	target_n
	)

	return cosine + mse


	# =========================================================
	# EVALUATION
	# =========================================================

	@torch.no_grad()
	def evaluate(ids):

	A.eval()
	B.eval()

	cos_sum = 0.0
	mse_sum = 0.0
	count = 0

	for start in range(0, len(ids), BATCH_SIZE):

	idx = ids[start:start + BATCH_SIZE]

	x = sense[idx].float().to(device)
	y = h3[idx].float().to(device)

	pred = bridge_forward(x)

	cos = torch.nn.functional.cosine_similarity(
	pred,
	y,
	dim=-1
	)

	mse = torch.nn.functional.mse_loss(
	pred,
	y,
	reduction="none"
	).mean(dim=1)

	cos_sum += cos.sum().item()
	mse_sum += mse.sum().item()

	count += len(idx)

	return cos_sum / count, mse_sum / count


	# =========================================================
	# BASELINE
	# =========================================================

	print()
	print("=" * 80)
	print("TRAINING")
	print("=" * 80)


	# =========================================================
	# TRAIN
	# =========================================================

	for epoch in range(EPOCHS):

	A.train()
	B.train()

	shuffled = train_ids[
	torch.randperm(len(train_ids), generator=g)
	]

	running = 0.0
	batches = 0

	for start in range(0, len(shuffled), BATCH_SIZE):

	idx = shuffled[start:start + BATCH_SIZE]

	x = sense[idx].float().to(device)
	y = h3[idx].float().to(device)

	optimizer.zero_grad(set_to_none=True)

	pred = bridge_forward(x)

	loss = loss_fn(pred, y)

	loss.backward()

	torch.nn.utils.clip_grad_norm_(
	params,
	1.0
	)

	optimizer.step()

	running += loss.item()
	batches += 1

	if batches % 100 == 0:

	print(
	f"Epoch {epoch + 1}/{EPOCHS} \| "
	f"batch {batches} \| "
	f"loss {running / batches:.6f}"
	)

	val_cos, val_mse = evaluate(val_ids)

	print()
	print(
	f"EPOCH {epoch + 1} COMPLETE \| "
	f"train_loss={running / batches:.6f} \| "
	f"val_cos={val_cos:.6f} \| "
	f"val_mse={val_mse:.8f}"
	)
	print()


	# =========================================================
	# FINAL EVALUATION
	# =========================================================

	val_cos, val_mse = evaluate(val_ids)

	print("=" * 80)
	print("FINAL")
	print("=" * 80)

	print("Validation cosine similarity:", val_cos)
	print("Validation MSE:", val_mse)


	# =========================================================
	# SAVE
	# =========================================================

	save_file(
	{
	"proj_in.weight": A.weight.detach().cpu().to(torch.float16),
	"proj_out.weight": B.weight.detach().cpu().to(torch.float16),
	},
	OUT,
	metadata={
	"source": "SenseNova-U1.5-8B-MoT",
	"target": "MiniMax-H3-Qwen3VL32B",
	"input_dim": str(INPUT_DIM),
	"rank": str(RANK),
	"output_dim": str(OUTPUT_DIM),
	"validation_cosine": str(val_cos),
	"validation_mse": str(val_mse),
	}
	)

	print()
	print("Saved bridge:")
	print(OUT)
	print()
	print("DONE")