apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-26 17:11:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

autocast bfloat16

Browse source
This commit is contained in:
51616 2025-05-30 08:11:31 +00:00
parent f8a7ef4b92
commit 75fc4e5ff8
1 changed files with 3 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
src/ctx_to_lora/modeling/hypernet.py
View file

@ -437,7 +437,7 @@ class HyperLoRA(nn.Module):
# hidden_size=self.config.base_hidden_size, # hidden_size=self.config.base_hidden_size,
# ) # )
# @torch.autocast(device_type="cuda", dtype=torch.bfloat16) @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def _to_lora_dict( def _to_lora_dict(
self, flat_loras: Float[Tensor, "bs n_layers n_modules r max_io_dim"] self, flat_loras: Float[Tensor, "bs n_layers n_modules r max_io_dim"]
) -> dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]]: ) -> dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]]:
@ -491,7 +491,7 @@ class HyperLoRA(nn.Module):
) )
return {k: v for k, v in zip(self.extra_modules, layernorms)} return {k: v for k, v in zip(self.extra_modules, layernorms)}
# @torch.autocast(device_type="cuda", dtype=torch.bfloat16) @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def forward( def forward(
self, self,
features: Float[Tensor, "bs seq_len feature_dim"], features: Float[Tensor, "bs seq_len feature_dim"],
@ -521,7 +521,7 @@ class HyperLoRA(nn.Module):
return flat_loras, flat_layernorms return flat_loras, flat_layernorms
# @torch.autocast(device_type="cuda", dtype=torch.bfloat16) @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def generate_weights( def generate_weights(
self, self,
features: Float[Tensor, "bs seq_len feature_dim"], features: Float[Tensor, "bs seq_len feature_dim"],