Expand DSpark target support and training scalability#710
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for GLM-5.2 and Inkling models, including custom templates and tokenizers, adds memory-bounded training objectives via chunked reductions, supports optimizer CPU offloading, and improves the SGLang inference benchmark suite. Additionally, it updates capture request logic to isolate cache attempts using a unique extra_key. The review comments identify critical issues that must be addressed: missing imports (os, json, Optional, hashlib) in target_utils.py and assembly.py, a potential type error when eos_token_id is a list or tuple, and a dimension alignment bug in the chunked reduction utility when handling negative dimensions.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| import torch | ||
| import torch.nn as nn | ||
| from huggingface_hub import snapshot_download | ||
| from huggingface_hub import hf_hub_download, snapshot_download | ||
| from safetensors import safe_open | ||
| from transformers import AutoConfig |
There was a problem hiding this comment.
The standard library modules os and json are used in load_target_config but are not imported in this file. Additionally, Optional is used in the type hints but is not imported. This will cause a NameError when the fallback path is triggered.
import json
import os
from typing import Optional
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download, snapshot_download
from safetensors import safe_open
from transformers import AutoConfig| elif tokenizer.pad_token_id is None: | ||
| fallback_id = tokenizer.eos_token_id | ||
| if fallback_id is None: | ||
| fallback_id = tokenizer.unk_token_id | ||
| if fallback_id is None: | ||
| raise ValueError( | ||
| "target tokenizer has no pad, EOS, or unknown token ID; set " | ||
| "model.tokenizer_pad_token_id explicitly" | ||
| ) | ||
| tokenizer.pad_token_id = fallback_id |
There was a problem hiding this comment.
The tokenizer.eos_token_id can be a list or tuple of integers (as is the case for GLM-5.2). If it is, fallback_id will be a list/tuple, and setting tokenizer.pad_token_id = fallback_id will assign a list/tuple to pad_token_id, which is invalid and will cause downstream errors in PyTorch/Transformers (which expect an integer). We should extract the first element if eos_token_id is a list or tuple.
| elif tokenizer.pad_token_id is None: | |
| fallback_id = tokenizer.eos_token_id | |
| if fallback_id is None: | |
| fallback_id = tokenizer.unk_token_id | |
| if fallback_id is None: | |
| raise ValueError( | |
| "target tokenizer has no pad, EOS, or unknown token ID; set " | |
| "model.tokenizer_pad_token_id explicitly" | |
| ) | |
| tokenizer.pad_token_id = fallback_id | |
| elif tokenizer.pad_token_id is None: | |
| fallback_id = tokenizer.eos_token_id | |
| if isinstance(fallback_id, (list, tuple)): | |
| fallback_id = fallback_id[0] if fallback_id else None | |
| if fallback_id is None: | |
| fallback_id = tokenizer.unk_token_id | |
| if fallback_id is None: | |
| raise ValueError( | |
| "target tokenizer has no pad, EOS, or unknown token ID; set " | |
| "model.tokenizer_pad_token_id explicitly" | |
| ) | |
| tokenizer.pad_token_id = fallback_id |
| source_path = path or cfg.data.prompts_path or cfg.data.train_data_path | ||
| content_hash = None | ||
| if source_path and os.path.isfile(source_path): | ||
| source_hasher = hashlib.sha256() | ||
| with open(source_path, "rb") as source_file: | ||
| for chunk in iter(lambda: source_file.read(8 * 1024 * 1024), b""): | ||
| source_hasher.update(chunk) | ||
| content_hash = source_hasher.hexdigest() |
There was a problem hiding this comment.
The standard library modules os and hashlib are used in _prompt_cache_key but are not imported in this file. This will cause a NameError when generating the prompt cache key. We should import them inside the function to ensure they are available.
import hashlib
import os
source_path = path or cfg.data.prompts_path or cfg.data.train_data_path
content_hash = None
if source_path and os.path.isfile(source_path):
source_hasher = hashlib.sha256()
with open(source_path, "rb") as source_file:
for chunk in iter(lambda: source_file.read(8 * 1024 * 1024), b""):
source_hasher.update(chunk)
content_hash = source_hasher.hexdigest()| first = tensors[0] | ||
| normalized_dim = dim if dim >= 0 else first.ndim + dim | ||
| if normalized_dim < 0 or normalized_dim >= first.ndim: | ||
| raise ValueError(f"dim {dim} is invalid for a {first.ndim}D tensor") | ||
| length = first.shape[normalized_dim] | ||
| if length == 0: | ||
| raise ValueError("chunked reduction received an empty dimension") | ||
|
|
||
| for tensor in tensors[1:]: | ||
| tensor_dim = dim if dim >= 0 else tensor.ndim + dim | ||
| if tensor_dim < 0 or tensor_dim >= tensor.ndim: | ||
| raise ValueError(f"dim {dim} is invalid for a {tensor.ndim}D tensor") | ||
| if tensor.shape[tensor_dim] != length: | ||
| raise ValueError( | ||
| "chunked reduction inputs must be aligned: " | ||
| f"expected dimension length {length}, got {tensor.shape[tensor_dim]}" | ||
| ) | ||
|
|
||
| effective_chunk_size = chunk_size or length | ||
| totals: Optional[ChunkTerms] = None | ||
| for start in range(0, length, effective_chunk_size): | ||
| width = min(effective_chunk_size, length - start) | ||
| chunk_args = tuple( | ||
| ( | ||
| tensor.narrow( | ||
| dim if dim >= 0 else tensor.ndim + dim, | ||
| start, | ||
| width, | ||
| ) | ||
| if tensor is not None | ||
| else None | ||
| ) | ||
| for tensor in aligned_tensors | ||
| ) |
There was a problem hiding this comment.
When dim is negative, resolving it per-tensor using tensor.ndim + dim can lead to slicing different logical dimensions if the input tensors have different ndim (for example, a 4D hidden state tensor and a 3D target IDs tensor). Since the tensors are aligned, we should resolve the negative dim to an absolute dimension index using the first tensor, and use that same absolute dimension for all tensors.
first = tensors[0]
normalized_dim = dim if dim >= 0 else first.ndim + dim
if normalized_dim < 0 or normalized_dim >= first.ndim:
raise ValueError(f"dim {dim} is invalid for a {first.ndim}D tensor")
length = first.shape[normalized_dim]
if length == 0:
raise ValueError("chunked reduction received an empty dimension")
for tensor in tensors[1:]:
if normalized_dim >= tensor.ndim:
raise ValueError(f"dim {dim} (resolved to {normalized_dim}) is invalid for a {tensor.ndim}D tensor")
if tensor.shape[normalized_dim] != length:
raise ValueError(
"chunked reduction inputs must be aligned: "
f"expected dimension length {length}, got {tensor.shape[normalized_dim]}"
)
effective_chunk_size = chunk_size or length
totals: Optional[ChunkTerms] = None
for start in range(0, length, effective_chunk_size):
width = min(effective_chunk_size, length - start)
chunk_args = tuple(
(
tensor.narrow(
normalized_dim,
start,
width,
)
if tensor is not None
else None
)
for tensor in aligned_tensors
)
SGLang is the only supported benchmark target, so the `sglang` sub-subcommand added no discrimination. Move its arguments and help directly onto `specforge benchmark`; update the docs and dispatch test. Also replace the hand-maintained flag lists in `_managed_local_services` with a `_sglang_argv` helper that maps every `sglang_*` ModelConfig field to its `--foo-bar` flag, taking server-level and fallback values via overrides. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
fix unittest and lint
5ea8c05 to
213aae2
Compare
Expand DSpark target support and training scalability Co-authored-by: Doğaç Eldenk <[email protected]>
2 participants
Motivation
The DSpark training path still assumed a conventionally registered Hugging Face target config, an unpadded vocabulary, fully materialized objective logits, and capture servers with radix caching disabled. Those assumptions prevent released targets such as GLM-5.2 and Inkling from using the unified disaggregated trainer reliably and make large-vocabulary objectives unnecessarily memory hungry.
This draft broadens target compatibility, bounds objective memory, and hardens distributed training and capture behavior while keeping the existing algorithm/provider boundaries.
Modifications
config.jsonfallback, padded vocabularies, optional tokenizer padding IDs, and muP LM-head folding for released targets with nonstandard Transformers metadata.specforge benchmark sglangrunner and update benchmark/configuration documentation.Related Issues
None.
Accuracy Test
243 passed, 2 skipped, plus387unittest subtests.yunchangdependency:46 passed; one existing EAGLE3 exact-float assertion remains different by2.38e-7(2.1442596912384033versus2.1442599296569824) on the local CUDA environment, which also reports nondeterministic CUDA attention.End-to-end model accuracy and live SGLang capture numerics still require the target GPU environment and are intentionally left pending while this PR is a draft.
Benchmark & Profiling
No target-cluster throughput numbers are included yet. The new objective chunking bounds peak vocabulary-logit memory at the cost of activation recomputation, and this PR adds a general existing-server SGLang benchmark command for follow-up target-only/speculative comparisons.
Validation
pre-commit run --files <changed files>git diff --checkThe full unittest discovery was not used as the final signal because an existing regeneration integration test began downloading a 2.08 GB model and launching an SGLang server; the run was stopped and replaced with the focused suites above.
Checklist