[Model] H2O Danube3-4b (#6451)

This commit is contained in:
Joe
2024-07-26 20:47:50 -07:00
committed by GitHub
parent ed94e4f427
commit 14dbd5a767
10 changed files with 79 additions and 7 deletions

View File

@@ -508,6 +508,12 @@ def create_kv_caches_with_random(
seed: int = 0,
device: Optional[str] = "cuda",
) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
if cache_dtype == "fp8" and head_size % 16:
raise ValueError(
f"Does not support key cache of type fp8 with head_size {head_size}"
)
torch.random.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)