ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(spark): whole-token device-name match so GB10 != GB100/GB10X

Browse source 
A loose substring match ("GB10" in name) misdetected a discrete Grace+Blackwell
datacenter GPU (e.g. nvidia-smi name containing "GB100") as a unified-memory DGX
Spark, applying the UMA tuning (pin_memory off, vLLM disabled, allocator capped to
0.80) and regressing that hardware. Match each device-name token with non-alphanumeric
boundaries instead. Found by a platform x device-name gating simulation; the real N1X
(JMJWOA-Generic-GPU) still detects, GB100/B100/GB200/GH200/B200 now correctly reject.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Daniel Han 2026-06-11 22:21:46 -07:00
commit 4cebfabaa6
3 changed files with 24 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

8
studio/backend/core/training/worker.py
View file

@ -749,8 +749,10 @@ def _nvidia_classify_spark_unified_memory(props: Any) -> tuple[str, bool]:
if getattr(props, "is_integrated", 0):
return "is_integrated", True
name_upper = (getattr(props, "name", "") or "").upper()
import re
for token in ("GB10", "GB110", "JMJWOA", "N1X", "DGX SPARK"):
if token in name_upper:
# Whole-token match so "GB10" does not match a discrete "GB100"/"GB10X".
if re.search(r"(?<![A-Z0-9])" + re.escape(token) + r"(?![A-Z0-9])", name_upper):
return token, True
return "", False
@ -2455,8 +2457,10 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
timeout = 5,
)
_names_u = (_smi.stdout or "").upper()
import re as _re
_spark_smi = any(
t in _names_u for t in ("GB10", "GB110", "JMJWOA", "N1X", "DGX SPARK")
_re.search(r"(?<![A-Z0-9])" + _re.escape(t) + r"(?![A-Z0-9])", _names_u)
for t in ("GB10", "GB110", "JMJWOA", "N1X", "DGX SPARK")
)
if _spark_smi and os.environ.get("UNSLOTH_NO_EXPANDABLE_SEGMENTS") != "1":
_conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")

7
unsloth/kernels/flex_attention.py
View file

@ -49,7 +49,12 @@ def _flex_is_dgx_spark():
timeout = 5,
)
names = (out.stdout or "").upper()
return any(t in names for t in ("GB10", "JMJWOA", "N1X", "DGX SPARK", "GB110"))
# Whole-token match so "GB10" does not match a discrete "GB100"/"GB10X".
import re
return any(
re.search(r"(?<![A-Z0-9])" + re.escape(t) + r"(?![A-Z0-9])", names)
for t in ("GB10", "JMJWOA", "N1X", "DGX SPARK", "GB110")
)
except Exception:
return False

14
unsloth/models/_utils.py
View file

@ -998,6 +998,16 @@ from transformers.modeling_utils import logger as transformers_logger
_DGX_SPARK_DEVICE_TOKENS = ("GB10", "JMJWOA", "N1X", "DGX SPARK", "GB110")
def _name_has_spark_token(names_upper):
# Whole-token match so "GB10" does NOT match "GB100"/"GB10X" -- a discrete
# Grace+Blackwell datacenter GPU must not be misread as a unified-memory Spark.
import re
return any(
re.search(r"(?<![A-Z0-9])" + re.escape(tok) + r"(?![A-Z0-9])", names_upper)
for tok in _DGX_SPARK_DEVICE_TOKENS
)
@functools.lru_cache(maxsize = None)
def is_dgx_spark():
"""True only on DGX Spark / N1X Spark-class machines (gate: aarch64 + NVIDIA
@ -1017,7 +1027,7 @@ def is_dgx_spark():
names = " ".join(
str(torch.cuda.get_device_name(i)).upper() for i in range(torch.cuda.device_count())
)
return any(token in names for token in _DGX_SPARK_DEVICE_TOKENS)
return _name_has_spark_token(names)
except Exception:
return False
@ -1046,7 +1056,7 @@ def _is_dgx_spark_no_cuda_init():
timeout = 5,
)
names = (out.stdout or "").upper()
return any(token in names for token in _DGX_SPARK_DEVICE_TOKENS)
return _name_has_spark_token(names)
except Exception:
return False