Add distributed server implementation and protocol definitions
- Introduced `protocol.py` for shared data models used in server/client communication, including request and response schemas for registration, job submission, and status updates. - Implemented `server.py` to manage a SQLite job queue and client registry, handling job polling, status updates, and job completion. - Created a cheat sheet for server usage, detailing commands for starting the server, submitting jobs, and monitoring clients. - Added several experiment configuration files for various training setups, including geometry vector injections and baseline ensembles.
This commit is contained in:
@@ -0,0 +1,46 @@
|
||||
"""mono_bridge — MonoBridge: passthrough for single-tower fusion stages.
|
||||
|
||||
The v4 stage runner always expects tower → bridge → head. For configs that
|
||||
have only one tower feeding a head, MonoBridge is the no-op bridge that lets
|
||||
the architecture be "head sits directly on tower" without any extra projection,
|
||||
SE, or fusion logic.
|
||||
|
||||
Optional LayerNorm is exposed for consistency with FusionBridge but defaults
|
||||
off to keep the embedding numerically identical to the tower's output.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class MonoBridge(nn.Module):
|
||||
"""Single-input passthrough bridge.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input_dims : list[int] — must be length 1
|
||||
use_ln : if True, wrap the embedding in a LayerNorm
|
||||
"""
|
||||
|
||||
def __init__(self, input_dims: list[int], use_ln: bool = False):
|
||||
super().__init__()
|
||||
if len(input_dims) != 1:
|
||||
raise ValueError(
|
||||
f"MonoBridge expects exactly 1 input dim, got {len(input_dims)}"
|
||||
)
|
||||
self.out_dim = input_dims[0]
|
||||
self.ln = nn.LayerNorm(self.out_dim) if use_ln else nn.Identity()
|
||||
|
||||
def forward(self, embeddings: list[torch.Tensor]) -> torch.Tensor:
|
||||
if len(embeddings) != 1:
|
||||
raise ValueError(
|
||||
f"MonoBridge forward expects 1 embedding, got {len(embeddings)}"
|
||||
)
|
||||
return self.ln(embeddings[0])
|
||||
|
||||
def set_phase(self, phase: str) -> None:
|
||||
"""Freeze during tower_warmup; trainable otherwise (matches FusionBridge)."""
|
||||
enabled = phase not in ("tower_warmup", "cd_warmup")
|
||||
for p in self.parameters():
|
||||
p.requires_grad_(enabled)
|
||||
Reference in New Issue
Block a user