diff --git a/easycom_head.py b/easycom_head.py index 59cd8cb..0802c7c 100644 --- a/easycom_head.py +++ b/easycom_head.py @@ -46,12 +46,32 @@ class TurnPredictionHead(nn.Module): num_horizons: int = NUM_HORIZONS, num_classes: int = NUM_CLASSES, dropout: float = 0.1, + num_input_layers: int = 1, + learnable_layer_mixture: bool = False, ) -> None: super().__init__() + if num_input_layers <= 0: + raise ValueError("num_input_layers must be positive") + if learnable_layer_mixture and num_input_layers < 2: + raise ValueError( + "learnable_layer_mixture requires at least two input layers" + ) + if not learnable_layer_mixture and num_input_layers != 1: + raise ValueError( + "Multiple input layers require learnable_layer_mixture=True" + ) self.hidden_size = hidden_size self.num_horizons = num_horizons self.num_classes = num_classes self.dropout = dropout + self.num_input_layers = num_input_layers + self.learnable_layer_mixture = learnable_layer_mixture + if learnable_layer_mixture: + # Zero initialization exactly reproduces a uniform layer mean at + # step zero while allowing training to select useful query layers. + self.layer_logits = nn.Parameter(torch.zeros(num_input_layers)) + else: + self.register_parameter("layer_logits", None) self.net = nn.Sequential( nn.LayerNorm(hidden_size), nn.Linear(hidden_size, hidden_size // 2), @@ -61,11 +81,36 @@ class TurnPredictionHead(nn.Module): ) def forward(self, states: torch.Tensor) -> torch.Tensor: + if states.shape[-1] != self.hidden_size: + raise ValueError( + f"Expected hidden size {self.hidden_size}, got {states.shape[-1]}" + ) + if self.learnable_layer_mixture: + if states.ndim < 3 or states.shape[-2] != self.num_input_layers: + raise ValueError( + "Layer-mixture input must end in " + f"[{self.num_input_layers},{self.hidden_size}], got " + f"{tuple(states.shape)}" + ) + assert self.layer_logits is not None + weights = self.layer_logits.softmax(dim=0).to(states.dtype) + view_shape = (1,) * (states.ndim - 2) + ( + self.num_input_layers, + 1, + ) + states = (states * weights.view(view_shape)).sum(dim=-2) logits = self.net(states) return logits.view( *states.shape[:-1], self.num_horizons, self.num_classes ) + def layer_mixture_weights(self) -> torch.Tensor | None: + """Return normalized oldest-to-newest layer weights, if enabled.""" + + if self.layer_logits is None: + return None + return self.layer_logits.softmax(dim=0) + def _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> None: if logits.shape != targets.shape: @@ -219,6 +264,10 @@ class SessionData: def hidden_size(self) -> int: return int(self.states.shape[-1]) + @property + def num_input_layers(self) -> int: + return 1 if self.states.ndim == 2 else int(self.states.shape[-2]) + def future_examples( self, num_horizons: int = NUM_HORIZONS, @@ -256,8 +305,12 @@ def load_session_data( if not isinstance(payload, dict) or "states" not in payload: raise ValueError(f"Invalid feature payload: {feature_path}") states = payload["states"] - if not isinstance(states, torch.Tensor) or states.ndim != 2: - raise ValueError(f"states must be [T,D] in {feature_path}") + if not isinstance(states, torch.Tensor) or states.ndim not in {2, 3}: + raise ValueError(f"states must be [T,D] or [T,L,D] in {feature_path}") + if states.ndim == 3 and states.shape[1] < 2: + raise ValueError( + f"Stacked states must contain at least two layers: {feature_path}" + ) if not torch.is_floating_point(states): raise ValueError(f"states must be floating point in {feature_path}") if not torch.isfinite(states).all(): @@ -351,6 +404,8 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]: config.setdefault("num_horizons", NUM_HORIZONS) config.setdefault("num_classes", NUM_CLASSES) config.setdefault("dropout", 0.1) + config.setdefault("num_input_layers", 1) + config.setdefault("learnable_layer_mixture", False) return config @@ -360,4 +415,4 @@ def validate_disjoint_sessions( ) -> None: overlap = set(first) & set(second) if overlap: - raise ValueError(f"Session sets overlap: {sorted(overlap)}") \ No newline at end of file + raise ValueError(f"Session sets overlap: {sorted(overlap)}") diff --git a/infer_turn_head.py b/infer_turn_head.py index 1f400e0..582c9ee 100644 --- a/infer_turn_head.py +++ b/infer_turn_head.py @@ -283,6 +283,11 @@ def main() -> None: f"{name} feature dim is {data.hidden_size}, checkpoint expects " f"{head_config['hidden_size']}" ) + if data.num_input_layers != int(head_config["num_input_layers"]): + raise ValueError( + f"{name} has {data.num_input_layers} feature layers, checkpoint " + f"expects {head_config['num_input_layers']}" + ) probabilities = _predict_states( head, data.states, @@ -323,14 +328,29 @@ def main() -> None: "checkpoint_epoch": checkpoint_payload.get("epoch"), "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"), "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"), + "layer_mixture_weights_oldest_to_newest": ( + head.layer_mixture_weights().detach().float().cpu().tolist() + if head.layer_mixture_weights() is not None + else None + ), } ) metrics_path = args.output_dir / "metrics.json" metrics_path.write_text(json.dumps(metrics, indent=2) + "\n") _print_metrics(metrics) + if metrics["layer_mixture_weights_oldest_to_newest"] is not None: + print( + "\nLayer mixture weights (oldest->newest): " + + ", ".join( + f"{value:.6f}" + for value in metrics[ + "layer_mixture_weights_oldest_to_newest" + ] + ) + ) print(f"\nmetrics: {metrics_path}") print(f"EgoSpeak-compatible predictions: {prediction_dir}") if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/run_extract_easycom_direct.sh b/run_extract_easycom_direct.sh index ec7aff4..697436f 100644 --- a/run_extract_easycom_direct.sh +++ b/run_extract_easycom_direct.sh @@ -17,6 +17,7 @@ MEMORY_SECONDS="${MEMORY_SECONDS:-20}" NATIVE_BLOCK_SECONDS="${NATIVE_BLOCK_SECONDS:-2}" VIDEO_FPS="${VIDEO_FPS:-2}" READOUT_LAST_N_LAYERS="${READOUT_LAST_N_LAYERS:-1}" +READOUT_LAYER_FUSION="${READOUT_LAYER_FUSION:-mean}" MODE="${MODE:-open-stream-replay}" DTYPE="${DTYPE:-bfloat16}" DECODER="${DECODER:-decord}" @@ -74,6 +75,7 @@ run_one() { echo "output=${output}" echo "memory_seconds=${MEMORY_SECONDS} video_fps=${VIDEO_FPS} mode=${MODE}" echo "readout_last_n_layers=${READOUT_LAST_N_LAYERS}" + echo "readout_layer_fusion=${READOUT_LAYER_FUSION}" echo "decoder=${DECODER} decoder_gpu_id=${DECODER_GPU_ID} frame_cache_size=${FRAME_CACHE_SIZE}" if ! CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \ @@ -87,6 +89,7 @@ run_one() { --native-block-seconds "${NATIVE_BLOCK_SECONDS}" \ --video-fps "${VIDEO_FPS}" \ --readout-last-n-layers "${READOUT_LAST_N_LAYERS}" \ + --readout-layer-fusion "${READOUT_LAYER_FUSION}" \ --frame-cache-size "${FRAME_CACHE_SIZE}" \ --decoder "${DECODER}" \ --decoder-gpu-id "${DECODER_GPU_ID}" \ diff --git a/streaming_turn_model.py b/streaming_turn_model.py index c7d26fb..f42b604 100644 --- a/streaming_turn_model.py +++ b/streaming_turn_model.py @@ -43,6 +43,7 @@ from typing import Any, Literal import torch from easycom_head import ( TurnPredictionHead, + checkpoint_head_config, egospeak_multihorizon_loss, miniroad_multihorizon_loss, multihorizon_cross_entropy, @@ -141,6 +142,7 @@ class StreamConfig: video_min_pixels: int = 64 * 28 * 28 video_max_pixels: int = 128 * 28 * 28 readout_last_n_layers: int = 1 + readout_layer_fusion: str = "mean" def validate(self) -> None: if self.step_seconds <= 0: @@ -157,6 +159,15 @@ class StreamConfig: raise ValueError("video_max_pixels must be >= video_min_pixels") if self.readout_last_n_layers <= 0: raise ValueError("readout_last_n_layers must be positive") + if self.readout_layer_fusion not in {"mean", "stack"}: + raise ValueError("readout_layer_fusion must be 'mean' or 'stack'") + if ( + self.readout_layer_fusion == "stack" + and self.readout_last_n_layers < 2 + ): + raise ValueError( + "stack readout requires readout_last_n_layers >= 2" + ) _require_integer_ratio( self.native_block_seconds, self.step_seconds, "native block / step" ) @@ -170,7 +181,7 @@ class ExtractionResult: """One fused Thinker state for every causal prediction time.""" times: torch.Tensor # [N], seconds in the original session - states: torch.Tensor # [N, D] + states: torch.Tensor # [N,D] or preserved query layers [N,L,D] verification: list[dict[str, Any]] @@ -956,6 +967,7 @@ def _read_state( next_rope_offset: int, query_ids: torch.LongTensor, last_n_layers: int = 1, + layer_fusion: str = "mean", ) -> torch.Tensor: """Read one q_t without committing the artificial assistant query.""" @@ -989,25 +1001,38 @@ def _read_state( return_dict=True, cache_position=cache_position, ) - return _readout_hidden_state(outputs, last_n_layers=last_n_layers) + return _readout_hidden_state( + outputs, + last_n_layers=last_n_layers, + layer_fusion=layer_fusion, + ) -def _readout_hidden_state(outputs: Any, *, last_n_layers: int) -> torch.Tensor: - """Return the query state, optionally averaging the final decoder layers.""" +def _readout_hidden_state( + outputs: Any, + *, + last_n_layers: int, + layer_fusion: str, +) -> torch.Tensor: + """Return one query vector or an ordered stack of final-layer vectors.""" + if layer_fusion not in {"mean", "stack"}: + raise ValueError(f"Unknown readout layer fusion: {layer_fusion}") if last_n_layers == 1: - return outputs.last_hidden_state[:, -1, :] - hidden_states = outputs.hidden_states - if hidden_states is None or len(hidden_states) < last_n_layers: - available = 0 if hidden_states is None else len(hidden_states) - raise RuntimeError( - f"Requested the last {last_n_layers} hidden states, but the model " - f"returned {available}" + stacked = outputs.last_hidden_state[:, -1:, :] + else: + hidden_states = outputs.hidden_states + if hidden_states is None or len(hidden_states) < last_n_layers: + available = 0 if hidden_states is None else len(hidden_states) + raise RuntimeError( + f"Requested the last {last_n_layers} hidden states, but the model " + f"returned {available}" + ) + stacked = torch.stack( + [state[:, -1, :] for state in hidden_states[-last_n_layers:]], + dim=1, ) - return torch.stack( - [state[:, -1, :] for state in hidden_states[-last_n_layers:]], - dim=0, - ).mean(dim=0) + return stacked.mean(dim=1) if layer_fusion == "mean" else stacked @torch.inference_mode() @@ -1019,6 +1044,7 @@ def _read_open_state( closure_position_ids: torch.LongTensor, query_ids: torch.LongTensor, last_n_layers: int = 1, + layer_fusion: str = "mean", ) -> torch.Tensor: """Temporarily close an open media turn, read q_t, and discard the branch.""" @@ -1078,7 +1104,11 @@ def _read_open_state( return_dict=True, cache_position=cache_position, ) - return _readout_hidden_state(outputs, last_n_layers=last_n_layers) + return _readout_hidden_state( + outputs, + last_n_layers=last_n_layers, + layer_fusion=layer_fusion, + ) @torch.inference_mode() @@ -1195,6 +1225,7 @@ class CausalPrefixExtractor: next_rope_offset=next_offset, query_ids=self.query_ids, last_n_layers=self.config.readout_last_n_layers, + layer_fusion=self.config.readout_layer_fusion, ) def extract( @@ -1358,6 +1389,7 @@ class OpenStreamReplayExtractor: closure_position_ids=closure_position_ids, query_ids=self.query_ids, last_n_layers=self.config.readout_last_n_layers, + layer_fusion=self.config.readout_layer_fusion, ) def _cache_from_intervals( @@ -1567,6 +1599,7 @@ class OpenStreamReplayExtractor: closure_position_ids=closure_position_ids, query_ids=self.query_ids, last_n_layers=self.config.readout_last_n_layers, + layer_fusion=self.config.readout_layer_fusion, ) states.append(state.cpu()) @@ -1808,6 +1841,7 @@ class OverlappingWorkingMemoryExtractor: closure_position_ids=closure_position_ids, query_ids=self.query_ids, last_n_layers=self.config.readout_last_n_layers, + layer_fusion=self.config.readout_layer_fusion, ) def extract( @@ -2002,12 +2036,16 @@ def _compare_states( ) -> dict[str, Any]: cached_f = cached.float() exact_f = exact.float() + cached_flat = cached_f.reshape(1, -1) + exact_flat = exact_f.reshape(1, -1) return { "kind": kind, "time": float(current_time), "max_abs": float((cached_f - exact_f).abs().max().item()), "mean_abs": float((cached_f - exact_f).abs().mean().item()), - "cosine": float(nn.functional.cosine_similarity(cached_f, exact_f).item()), + "cosine": float( + nn.functional.cosine_similarity(cached_flat, exact_flat, dim=-1).item() + ), } @@ -2104,9 +2142,18 @@ def _load_head( hidden_size: int, device: torch.device, ) -> TurnPredictionHead: - head = TurnPredictionHead(hidden_size=hidden_size) payload = torch.load(checkpoint, map_location="cpu", weights_only=True) - state_dict = payload.get("head", payload) if isinstance(payload, dict) else payload + if isinstance(payload, dict) and "head" in payload: + config = checkpoint_head_config(payload) + if int(config["hidden_size"]) != hidden_size: + raise ValueError( + f"Head hidden size {config['hidden_size']} != model {hidden_size}" + ) + head = TurnPredictionHead(**config) + state_dict = payload["head"] + else: + head = TurnPredictionHead(hidden_size=hidden_size) + state_dict = payload head.load_state_dict(state_dict) return head.to(device).eval() @@ -2160,6 +2207,15 @@ def parse_args() -> argparse.Namespace: "causal readout token (default: 1 preserves the original result)." ), ) + parser.add_argument( + "--readout-layer-fusion", + choices=("mean", "stack"), + default="mean", + help=( + "Fuse selected query layers by a fixed mean, or preserve [N,D] " + "for a trainable downstream layer mixture." + ), + ) parser.add_argument("--verify-every", type=int, default=0) parser.add_argument( "--verify-continuous-every", @@ -2228,6 +2284,7 @@ def main() -> None: video_min_pixels=args.video_min_pixels, video_max_pixels=args.video_max_pixels, readout_last_n_layers=args.readout_last_n_layers, + readout_layer_fusion=args.readout_layer_fusion, ) config.validate() diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py index 0acff60..23f2ad6 100644 --- a/streaming_turn_model_directly.py +++ b/streaming_turn_model_directly.py @@ -1061,6 +1061,15 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: "causal readout token (default: 1 preserves the original result)." ), ) + parser.add_argument( + "--readout-layer-fusion", + choices=("mean", "stack"), + default="mean", + help=( + "Fuse selected query layers by a fixed mean, or preserve [N,D] " + "for a trainable downstream layer mixture." + ), + ) parser.add_argument("--frame-cache-size", type=int, default=8) parser.add_argument( "--audio-feature-device", @@ -1341,6 +1350,7 @@ def main( video_min_pixels=args.video_min_pixels, video_max_pixels=args.video_max_pixels, readout_last_n_layers=args.readout_last_n_layers, + readout_layer_fusion=args.readout_layer_fusion, ) config.validate() diff --git a/train_omni_turn_lora.py b/train_omni_turn_lora.py index ea8cfc7..3232591 100644 --- a/train_omni_turn_lora.py +++ b/train_omni_turn_lora.py @@ -726,12 +726,7 @@ def _load_initial_head( raise ValueError( f"Head hidden size {config['hidden_size']} != cache hidden size {hidden_size}" ) - head = TurnPredictionHead( - hidden_size=hidden_size, - num_horizons=int(config["num_horizons"]), - num_classes=int(config["num_classes"]), - dropout=float(config["dropout"]), - ) + head = TurnPredictionHead(**config) head.load_state_dict(payload.get("head", payload)) if verbose: print(f"Initialized classification head from {checkpoint}") diff --git a/train_turn_head.py b/train_turn_head.py index b84c47c..6e472d1 100644 --- a/train_turn_head.py +++ b/train_turn_head.py @@ -91,6 +91,14 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--warmup-epochs", type=float, default=10.0) parser.add_argument("--warmup-factor", type=float, default=0.3) parser.add_argument("--dropout", type=float, default=0.2) + parser.add_argument( + "--learnable-layer-mixture", + action="store_true", + help=( + "Learn softmax weights over stacked [T,L,D] query features before " + "the anticipation head." + ), + ) parser.add_argument( "--loss", choices=("lstr-mce", "miniroad-normalized"), @@ -151,11 +159,12 @@ def _load_examples( expected_mode: str | None, dataset: str, selected_indices: dict[str, set[int]] | None = None, -) -> tuple[torch.Tensor, torch.Tensor, int, list[dict[str, Any]]]: +) -> tuple[torch.Tensor, torch.Tensor, int, int, list[dict[str, Any]]]: all_states: list[torch.Tensor] = [] all_targets: list[torch.Tensor] = [] metadata: list[dict[str, Any]] = [] hidden_size: int | None = None + num_input_layers: int | None = None for name in sessions: data = load_session_data( feature_root, @@ -186,6 +195,13 @@ def _load_examples( f"Feature dimensions differ: {name} has {data.hidden_size}, " f"expected {hidden_size}" ) + if num_input_layers is None: + num_input_layers = data.num_input_layers + elif data.num_input_layers != num_input_layers: + raise ValueError( + f"Feature layer counts differ: {name} has " + f"{data.num_input_layers}, expected {num_input_layers}" + ) all_states.append(states) all_targets.append(targets) metadata.append( @@ -207,12 +223,13 @@ def _load_examples( f"valid_10_horizon_examples={dense_example_count}, " f"selected_examples={len(states)}" ) - if hidden_size is None: + if hidden_size is None or num_input_layers is None: raise ValueError("No sessions selected") return ( torch.cat(all_states, dim=0), torch.cat(all_targets, dim=0), hidden_size, + num_input_layers, metadata, ) @@ -318,6 +335,7 @@ def _checkpoint( train_sessions: list[str], dev_sessions: list[str], hidden_size: int, + num_input_layers: int, feature_metadata: list[dict[str, Any]], history: list[dict[str, Any]], ) -> dict[str, Any]: @@ -329,6 +347,8 @@ def _checkpoint( "num_horizons": NUM_HORIZONS, "num_classes": NUM_CLASSES, "dropout": args.dropout, + "num_input_layers": num_input_layers, + "learnable_layer_mixture": args.learnable_layer_mixture, }, "alignment": { "state_label_offset": 0, @@ -363,6 +383,7 @@ def _checkpoint( "scheduler": "linear_warmup_times_cosine", "seed": args.seed, "expected_mode": args.expected_mode, + "learnable_layer_mixture": args.learnable_layer_mixture, "supervision_policy": ( "transition-sparse-manifest" if args.sparse_anchor_cache_dir @@ -440,7 +461,7 @@ def main() -> None: print(f"device: {device}") print(f"train sessions: {', '.join(train_sessions)}") print(f"dev sessions: {', '.join(dev_sessions) if dev_sessions else 'none'}") - states, targets, hidden_size, feature_metadata = _load_examples( + states, targets, hidden_size, num_input_layers, feature_metadata = _load_examples( args.feature_root, args.target_root, train_sessions, @@ -448,11 +469,29 @@ def main() -> None: args.dataset, sparse_anchor_indices, ) - print(f"training examples: {len(states)}, feature_dim={hidden_size}") + print( + f"training examples: {len(states)}, feature_dim={hidden_size}, " + f"input_layers={num_input_layers}" + ) + if num_input_layers > 1 and not args.learnable_layer_mixture: + raise ValueError( + f"Features contain {num_input_layers} stacked layers; pass " + "--learnable-layer-mixture" + ) + if num_input_layers == 1 and args.learnable_layer_mixture: + raise ValueError( + "--learnable-layer-mixture requires stacked [T,L,D] features" + ) dev_states: torch.Tensor | None = None dev_targets: torch.Tensor | None = None if dev_sessions: - dev_states, dev_targets, dev_hidden_size, _ = _load_examples( + ( + dev_states, + dev_targets, + dev_hidden_size, + dev_num_input_layers, + _, + ) = _load_examples( args.feature_root, args.target_root, dev_sessions, @@ -465,6 +504,11 @@ def main() -> None: f"Dev feature dimension {dev_hidden_size} differs from " f"training dimension {hidden_size}" ) + if dev_num_input_layers != num_input_layers: + raise ValueError( + f"Dev feature layer count {dev_num_input_layers} differs from " + f"training count {num_input_layers}" + ) print(f"development examples: {len(dev_states)}") dataset = TensorDataset(states, targets) @@ -483,9 +527,30 @@ def main() -> None: num_horizons=NUM_HORIZONS, num_classes=NUM_CLASSES, dropout=args.dropout, + num_input_layers=num_input_layers, + learnable_layer_mixture=args.learnable_layer_mixture, ).to(device) + if head.layer_logits is None: + optimizer_parameters: Any = head.parameters() + else: + optimizer_parameters = [ + { + "params": [ + parameter + for name, parameter in head.named_parameters() + if name != "layer_logits" + ], + "weight_decay": args.weight_decay, + }, + { + "params": [head.layer_logits], + # Weight decay on softmax logits artificially pulls the layer + # distribution back toward uniform. + "weight_decay": 0.0, + }, + ] optimizer = torch.optim.Adam( - head.parameters(), + optimizer_parameters, lr=args.learning_rate, weight_decay=args.weight_decay, ) @@ -546,6 +611,21 @@ def main() -> None: "train_loss": train_loss, "learning_rate": scheduler.get_last_lr()[0], } + layer_weights = head.layer_mixture_weights() + if layer_weights is not None: + record["layer_mixture_weights_oldest_to_newest"] = ( + layer_weights.detach().float().cpu().tolist() + ) + weight_text = ( + "" + if layer_weights is None + else ", layer_weights=[" + + ",".join( + f"{value:.4f}" + for value in record["layer_mixture_weights_oldest_to_newest"] + ) + + "]" + ) if dev_sessions: assert dev_states is not None and dev_targets is not None @@ -560,7 +640,7 @@ def main() -> None: record["dev_valid_mean_map"] = dev_map print( f"epoch {epoch:02d}: train_loss={train_loss:.6f}, " - f"dev_valid_mean_mAP={dev_map * 100:.3f}" + f"dev_valid_mean_mAP={dev_map * 100:.3f}{weight_text}" ) if dev_map > float(best_dev_map): best_dev_map = dev_map @@ -576,6 +656,7 @@ def main() -> None: train_sessions=train_sessions, dev_sessions=dev_sessions, hidden_size=hidden_size, + num_input_layers=num_input_layers, feature_metadata=feature_metadata, history=history + [record], ), @@ -583,7 +664,9 @@ def main() -> None: ) print(f"saved new best checkpoint: {args.output}") else: - print(f"epoch {epoch:02d}: train_loss={train_loss:.6f}") + print( + f"epoch {epoch:02d}: train_loss={train_loss:.6f}{weight_text}" + ) history.append(record) if not dev_sessions: @@ -598,6 +681,7 @@ def main() -> None: train_sessions=train_sessions, dev_sessions=dev_sessions, hidden_size=hidden_size, + num_input_layers=num_input_layers, feature_metadata=feature_metadata, history=history, ), @@ -614,6 +698,13 @@ def main() -> None: summary = {key: value for key, value in payload.items() if key != "head"} summary_path.write_text(json.dumps(summary, indent=2) + "\n") print(f"training summary: {summary_path}") + if args.learnable_layer_mixture: + saved_logits = payload["head"]["layer_logits"].float() + saved_weights = saved_logits.softmax(dim=0).tolist() + print( + "saved layer mixture weights (oldest->newest): " + + ", ".join(f"{value:.6f}" for value in saved_weights) + ) if __name__ == "__main__":