tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2026-07-06 12:59:14 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge 7640a1781f into 1b1cc3c599

Browse Source
This commit is contained in:
RoomWithOutRoof 2026-03-15 02:57:40 +08:00 committed by GitHub
commit 71d90265f5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
scripts/base_train.py
View File

@ -218,12 +218,13 @@ def disable_fp8(model):
return return
# Swap Float8Linear -> Linear (our custom class that casts weights to match input dtype) # Swap Float8Linear -> Linear (our custom class that casts weights to match input dtype)
# Use device="meta" to avoid VRAM spike - the weight tensor will be swapped in afterwards
for parent, attr_name, fp8_module in fp8_locations: for parent, attr_name, fp8_module in fp8_locations:
linear = Linear( linear = Linear(
fp8_module.in_features, fp8_module.in_features,
fp8_module.out_features, fp8_module.out_features,
bias=fp8_module.bias is not None, bias=fp8_module.bias is not None,
device=fp8_module.weight.device, device="meta", # Use meta device to avoid unnecessary VRAM allocation
dtype=fp8_module.weight.dtype, dtype=fp8_module.weight.dtype,
) )
linear.weight = fp8_module.weight # share, don't copy linear.weight = fp8_module.weight # share, don't copy