(
self,
model: nn.Module,
router_path: str | Path,
config: Optional[TIDEConfig] = None,
use_cuda_kernels: bool = True,
)
| 63 | """ |
| 64 | |
| 65 | def __init__( |
| 66 | self, |
| 67 | model: nn.Module, |
| 68 | router_path: str | Path, |
| 69 | config: Optional[TIDEConfig] = None, |
| 70 | use_cuda_kernels: bool = True, |
| 71 | ): |
| 72 | super().__init__() |
| 73 | self.model = model |
| 74 | self.config = config or TIDEConfig() |
| 75 | self.adapter = get_adapter(model) |
| 76 | model_on_cuda = self._device.type == "cuda" |
| 77 | self.use_cuda = use_cuda_kernels and _cuda_available and model_on_cuda |
| 78 | |
| 79 | # Load routers |
| 80 | checkpoint = RouterCheckpoint.load(router_path, device=self._device_str) |
| 81 | self.routers: Dict[int, TokenRouter] = {} |
| 82 | for layer_idx, router in checkpoint.routers.items(): |
| 83 | router = router.to(self._device) |
| 84 | router.requires_grad_(False) |
| 85 | if self.use_cuda: |
| 86 | router.down.weight.data = router.down.weight.data.t().contiguous() |
| 87 | router.up.weight.data = router.up.weight.data.t().contiguous() |
| 88 | self.routers[layer_idx] = router |
| 89 | |
| 90 | self._layers = self.adapter.get_layers(model) |
| 91 | self._final_norm = self.adapter.get_final_norm(model) |
| 92 | self._lm_head = self.adapter.get_lm_head(model) |
| 93 | |
| 94 | self.last_stats: Optional[ExitStats] = None |
| 95 | logger.info( |
| 96 | f"TIDERuntime initialized: {len(self._layers)} layers, " |
| 97 | f"{len(self.routers)} routers, CUDA={'on' if self.use_cuda else 'off'}" |
| 98 | ) |
| 99 | |
| 100 | @property |
| 101 | def _device(self) -> torch.device: |
nothing calls this directly
no test coverage detected