[codex] Support Gemma4 QAT LoRA local training - #4
Draft
teilomillet wants to merge 8 commits into
Draft
Conversation
Reviewer's GuideAdds a Gemma4 text-only local training/sampling path with shared train/infer PEFT models, and introduces a new Scaleway remote backend (Terraform + FastAPI training server) plus an SROIE campaign and associated tooling/tests. Sequence diagram for Gemma4 text-only local samplingsequenceDiagram
participant LT as LocalTrainHelper
participant PE as PyTorchEngine
participant GM as gemma4_text
participant TM as Gemma4LanguageModel
LT->>PE: generate(prompt_ids_list, num_samples, max_tokens, temperature, top_p, compute_entropy)
PE->>GM: is_gemma4_text_model(model)
alt Gemma4_text_model
PE->>PE: _generate_gemma4_text(...)
PE->>GM: unwrap_peft_model(model)
GM-->>PE: unwrapped_model
PE->>TM: language_model(input_ids, attention_mask, past_key_values, use_cache=True, return_shared_kv_states=True)
TM-->>PE: last_hidden_state, past_key_values, shared_kv_states
PE->>TM: lm_head(last_hidden_state)
TM-->>PE: logits
PE->>PE: _sample_next_token(logits, temperature, top_p)
loop per_token_until_max_tokens_or_eos
PE->>GM: eos_token_ids(model)
GM-->>PE: eos_ids
PE->>PE: append token_ids, logprobs, entropies
end
PE-->>LT: list[SampleResult]
else Other_model
PE->>PE: standard_generate_path
end
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
Owner
Author
|
@sourcery-ai review |
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- The Gemma4-specific helpers for unwrapping PEFT models and detecting text-only Gemma4 models are duplicated between
pytorch_engine.pyandlocal_train_helper.py; consider centralizing these utilities in a shared module to avoid drift and keep the Gemma4 logic consistent. - In
ScalewayTrainHelper._sample_one, the handling oflogprobsand token IDs assumes atoken_idfield and falls back totokenor0; this will be brittle across different OpenAI-compatible servers, so it would be safer to consistently use the tokenizer to map returnedtokenstrings back to IDs and to define clearer behavior when logprobs/top_logprobs are missing. - The Terraform
caller_ipvariable defaults to0.0.0.0/0, effectively exposing ports 8000/8001 publicly; it may be safer to require an explicit non-default CIDR or at least log a clear warning when using the wide-open default.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The Gemma4-specific helpers for unwrapping PEFT models and detecting text-only Gemma4 models are duplicated between `pytorch_engine.py` and `local_train_helper.py`; consider centralizing these utilities in a shared module to avoid drift and keep the Gemma4 logic consistent.
- In `ScalewayTrainHelper._sample_one`, the handling of `logprobs` and token IDs assumes a `token_id` field and falls back to `token` or `0`; this will be brittle across different OpenAI-compatible servers, so it would be safer to consistently use the tokenizer to map returned `token` strings back to IDs and to define clearer behavior when logprobs/top_logprobs are missing.
- The Terraform `caller_ip` variable defaults to `0.0.0.0/0`, effectively exposing ports 8000/8001 publicly; it may be safer to require an explicit non-default CIDR or at least log a clear warning when using the wide-open default.
## Individual Comments
### Comment 1
<location path="retrain/inference_engine/pytorch_engine.py" line_range="74-77" />
<code_context>
"""Local PyTorch/PEFT inference engine."""
- def __init__(self, model_name, device, peft_config, dtype):
+ def __init__(self, model_name, device, peft_config, dtype, existing_model=None):
"""Load a PEFT-wrapped model for inference.
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Guard against device mismatch when reusing an existing model instance.
When `existing_model` is passed, we skip moving it to `device`, which is fine for `LocalTrainHelper` where `train_device == infer_device`, but risky elsewhere (e.g., `existing_model` on GPU and `device="cpu"`). Please either assert that `existing_model.device` matches `device`, or always move `existing_model` to `device` to avoid surprising, device-mismatch issues.
</issue_to_address>
### Comment 2
<location path="retrain/scaleway_backend.py" line_range="189-198" />
<code_context>
+ def _sample_one(
</code_context>
<issue_to_address>
**issue (bug_risk):** Handle vLLM logprobs payload more defensively and avoid mixing token text and IDs.
In `_sample_one`, this fallback can inject strings into `token_ids`:
```python
token_ids.append(entry["token_id"] if "token_id" in entry else entry.get("token", 0))
```
In OpenAI-style APIs `token` is usually a string, so downstream code expecting integer IDs may break. Prefer either (a) requiring `token_id` and raising if it’s missing, or (b) explicitly converting `token` to an ID via the tokenizer. Also, add a check that `lp_content` length matches the number of generated tokens so you fail fast instead of returning partial sequences if the remote payload shape changes.
</issue_to_address>
### Comment 3
<location path="retrain/scaleway_backend.py" line_range="75-84" />
<code_context>
+ def sample(
</code_context>
<issue_to_address>
**suggestion (performance):** Sampling issues one HTTP request per sample; consider using server-side batching.
This loops over `prompt_ids_list` and `num_samples`, calling `_sample_one` for each sample, which can cause many sequential HTTP calls and poor throughput for larger `batch_size × group_size`. Since the OpenAI-compatible API supports multiple `messages` and `n` in one call, consider a batched path that uses `n=num_samples` per prompt and/or sends multiple prompts per request where supported to reduce latency and GPU load.
Suggested implementation:
```python
def sample(
self,
prompt_ids_list: list[list[int]],
num_samples: int,
max_tokens: int,
temperature: float,
top_p: float,
) -> SampleBatch:
"""
Batched sampling: decodes token IDs to text once, then uses a batched
OpenAI-compatible request that leverages `n=num_samples` and multiple
prompts per request (where supported) to reduce HTTP overhead.
"""
from transformers import AutoTokenizer
# Lazy-load tokenizer for decoding token IDs back to text for the prompt
# vLLM OpenAI-compat API expects text prompts, not token IDs
tokenizer = AutoTokenizer.from_pretrained(self._model_name)
# Decode all prompts up-front. This keeps the public API in token-space
# (prompt_ids_list) while letting us use text prompts for the HTTP API.
decoded_prompts = [
tokenizer.decode(prompt_ids, skip_special_tokens=True)
for prompt_ids in prompt_ids_list
]
# Delegate to a batched helper that uses server-side batching where
# possible (multiple prompts per request and `n=num_samples`).
return self._sample_batched(
prompts=decoded_prompts,
num_samples=num_samples,
max_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
)
```
To fully implement server-side batching and avoid one HTTP call per sample, you should:
1. Implement a new helper method on this class (or reuse an existing one, if present) with a signature like:
```python
def _sample_batched(
self,
prompts: list[str],
num_samples: int,
max_tokens: int,
temperature: float,
top_p: float,
) -> SampleBatch:
```
This helper should:
- Use `self._client.post(...)` to call your OpenAI-compatible endpoint.
- Prefer a single request with:
- All prompts in one call if the endpoint supports batched prompts (e.g., `input=[...]` or similar).
- `n=num_samples` to request multiple samples per prompt server-side.
- Fall back to chunking `prompts` into smaller batches if the backend has a max batch size.
- Construct and return the `SampleBatch` object in the same format as the previous implementation (i.e., same shape and ordering as the old nested-loop `_sample_one` calls).
2. If `self._model_name` is not the correct identifier for `AutoTokenizer.from_pretrained(..)` in this codebase:
- Replace `self._model_name` with the appropriate configuration field (for example `self._config.model_name` or similar).
3. If the previous implementation of `sample` relied on a helper like `_sample_one(...)`:
- Move the HTTP call/conversion logic from `_sample_one` into `_sample_batched` so that:
- `_sample_batched` issues batched calls and builds the `SampleBatch`.
- Optionally, `_sample_one` becomes a thin wrapper that calls `_sample_batched` with `num_samples=1` for any remaining single-sample callers.
4. Ensure any existing callers of `sample(...)` still receive identical semantics:
- The total number of samples produced must be `len(prompt_ids_list) * num_samples`.
- The samples for each prompt should be grouped as before (e.g., first all samples for prompt 0, then for prompt 1, etc.), or adjust downstream code accordingly if you change this ordering.
</issue_to_address>
### Comment 4
<location path="retrain/scaleway/training_server.py" line_range="118-127" />
<code_context>
+# LoRA reload on inference engine
+# ---------------------------------------------------------------------------
+
+def _reload_lora_on_inference(name: str) -> None:
+ # The adapter was just saved by checkpoint(); we need to tell the inference
+ # engine to reload it. vLLM and SGLang use different endpoints.
+ try:
+ if _inference_engine == "vllm":
+ httpx.post(
+ f"{_inference_url}/v1/load_lora_adapter",
+ json={"lora_name": name, "lora_path": name},
+ timeout=30,
+ ).raise_for_status()
+ else:
+ httpx.post(
+ f"{_inference_url}/add_lora",
+ json={"lora_name": name, "lora_path": name},
+ timeout=30,
+ ).raise_for_status()
+ except Exception as exc:
+ logger.warning("LoRA reload on inference engine failed: %s", exc)
+
+
</code_context>
<issue_to_address>
**issue (bug_risk):** LoRA reload uses `lora_path=name`, which likely doesn’t match the actual adapter directory.
In `_reload_lora_on_inference`, you send `{"lora_name": name, "lora_path": name}` to vLLM and SGLang, but `checkpoint()` typically writes adapters under `adapter_path / name` (e.g. `/tmp/retrain_adapter/<name>`). The inference server likely needs that full path, not just `name`, and will fail to reload if it runs with a different CWD. Please either return the absolute adapter path from `checkpoint`/`save_adapter` and pass it into `_reload_lora_on_inference`, or define a shared base adapter directory and build the full path here.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+74
to
77
| def __init__(self, model_name, device, peft_config, dtype, existing_model=None): | ||
| """Load a PEFT-wrapped model for inference. | ||
|
|
||
| Args: |
There was a problem hiding this comment.
suggestion (bug_risk): Guard against device mismatch when reusing an existing model instance.
When existing_model is passed, we skip moving it to device, which is fine for LocalTrainHelper where train_device == infer_device, but risky elsewhere (e.g., existing_model on GPU and device="cpu"). Please either assert that existing_model.device matches device, or always move existing_model to device to avoid surprising, device-mismatch issues.
Comment on lines
+189
to
+198
| def _sample_one( | ||
| self, | ||
| prompt: str, | ||
| max_tokens: int, | ||
| temperature: float, | ||
| top_p: float, | ||
| ) -> tuple[list[int], list[float]]: | ||
| payload = { | ||
| "model": self._model, | ||
| "messages": [{"role": "user", "content": prompt}], |
There was a problem hiding this comment.
issue (bug_risk): Handle vLLM logprobs payload more defensively and avoid mixing token text and IDs.
In _sample_one, this fallback can inject strings into token_ids:
token_ids.append(entry["token_id"] if "token_id" in entry else entry.get("token", 0))In OpenAI-style APIs token is usually a string, so downstream code expecting integer IDs may break. Prefer either (a) requiring token_id and raising if it’s missing, or (b) explicitly converting token to an ID via the tokenizer. Also, add a check that lp_content length matches the number of generated tokens so you fail fast instead of returning partial sequences if the remote payload shape changes.
Comment on lines
+75
to
+84
| def sample( | ||
| self, | ||
| prompt_ids_list: list[list[int]], | ||
| num_samples: int, | ||
| max_tokens: int, | ||
| temperature: float, | ||
| top_p: float, | ||
| ) -> SampleBatch: | ||
| from transformers import AutoTokenizer | ||
| # Lazy-load tokenizer for decoding token IDs back to text for the prompt |
There was a problem hiding this comment.
suggestion (performance): Sampling issues one HTTP request per sample; consider using server-side batching.
This loops over prompt_ids_list and num_samples, calling _sample_one for each sample, which can cause many sequential HTTP calls and poor throughput for larger batch_size × group_size. Since the OpenAI-compatible API supports multiple messages and n in one call, consider a batched path that uses n=num_samples per prompt and/or sends multiple prompts per request where supported to reduce latency and GPU load.
Suggested implementation:
def sample(
self,
prompt_ids_list: list[list[int]],
num_samples: int,
max_tokens: int,
temperature: float,
top_p: float,
) -> SampleBatch:
"""
Batched sampling: decodes token IDs to text once, then uses a batched
OpenAI-compatible request that leverages `n=num_samples` and multiple
prompts per request (where supported) to reduce HTTP overhead.
"""
from transformers import AutoTokenizer
# Lazy-load tokenizer for decoding token IDs back to text for the prompt
# vLLM OpenAI-compat API expects text prompts, not token IDs
tokenizer = AutoTokenizer.from_pretrained(self._model_name)
# Decode all prompts up-front. This keeps the public API in token-space
# (prompt_ids_list) while letting us use text prompts for the HTTP API.
decoded_prompts = [
tokenizer.decode(prompt_ids, skip_special_tokens=True)
for prompt_ids in prompt_ids_list
]
# Delegate to a batched helper that uses server-side batching where
# possible (multiple prompts per request and `n=num_samples`).
return self._sample_batched(
prompts=decoded_prompts,
num_samples=num_samples,
max_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
)To fully implement server-side batching and avoid one HTTP call per sample, you should:
-
Implement a new helper method on this class (or reuse an existing one, if present) with a signature like:
def _sample_batched( self, prompts: list[str], num_samples: int, max_tokens: int, temperature: float, top_p: float, ) -> SampleBatch:
This helper should:
- Use
self._client.post(...)to call your OpenAI-compatible endpoint. - Prefer a single request with:
- All prompts in one call if the endpoint supports batched prompts (e.g.,
input=[...]or similar). n=num_samplesto request multiple samples per prompt server-side.
- All prompts in one call if the endpoint supports batched prompts (e.g.,
- Fall back to chunking
promptsinto smaller batches if the backend has a max batch size. - Construct and return the
SampleBatchobject in the same format as the previous implementation (i.e., same shape and ordering as the old nested-loop_sample_onecalls).
- Use
-
If
self._model_nameis not the correct identifier forAutoTokenizer.from_pretrained(..)in this codebase:- Replace
self._model_namewith the appropriate configuration field (for exampleself._config.model_nameor similar).
- Replace
-
If the previous implementation of
samplerelied on a helper like_sample_one(...):- Move the HTTP call/conversion logic from
_sample_oneinto_sample_batchedso that:_sample_batchedissues batched calls and builds theSampleBatch.- Optionally,
_sample_onebecomes a thin wrapper that calls_sample_batchedwithnum_samples=1for any remaining single-sample callers.
- Move the HTTP call/conversion logic from
-
Ensure any existing callers of
sample(...)still receive identical semantics:- The total number of samples produced must be
len(prompt_ids_list) * num_samples. - The samples for each prompt should be grouped as before (e.g., first all samples for prompt 0, then for prompt 1, etc.), or adjust downstream code accordingly if you change this ordering.
- The total number of samples produced must be
Comment on lines
+118
to
+127
| def _reload_lora_on_inference(name: str) -> None: | ||
| # The adapter was just saved by checkpoint(); we need to tell the inference | ||
| # engine to reload it. vLLM and SGLang use different endpoints. | ||
| try: | ||
| if _inference_engine == "vllm": | ||
| httpx.post( | ||
| f"{_inference_url}/v1/load_lora_adapter", | ||
| json={"lora_name": name, "lora_path": name}, | ||
| timeout=30, | ||
| ).raise_for_status() |
There was a problem hiding this comment.
issue (bug_risk): LoRA reload uses lora_path=name, which likely doesn’t match the actual adapter directory.
In _reload_lora_on_inference, you send {"lora_name": name, "lora_path": name} to vLLM and SGLang, but checkpoint() typically writes adapters under adapter_path / name (e.g. /tmp/retrain_adapter/<name>). The inference server likely needs that full path, not just name, and will fail to reload if it runs with a different CWD. Please either return the absolute adapter path from checkpoint/save_adapter and pass it into _reload_lora_on_inference, or define a shared base adapter directory and build the full path here.
Owner
Author
|
@sourcery-ai review |
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- In
ScalewayTrainHelper._wait_healthy, the samedeadlineis shared across both the inference and training health checks, which can make the effective timeout per service much lower thanhealth_timeout_s; consider using a separate deadline per URL or accounting for that explicitly. - In
TerraformRunner._run_tfyou raiseTerraformErrorwithout including stderr, which makes debugging failed applies/destroys harder; capturing stderr and including at least a truncated version in the exception message would improve diagnosability.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `ScalewayTrainHelper._wait_healthy`, the same `deadline` is shared across both the inference and training health checks, which can make the effective timeout per service much lower than `health_timeout_s`; consider using a separate deadline per URL or accounting for that explicitly.
- In `TerraformRunner._run_tf` you raise `TerraformError` without including stderr, which makes debugging failed applies/destroys harder; capturing stderr and including at least a truncated version in the exception message would improve diagnosability.
## Individual Comments
### Comment 1
<location path="retrain/inference_engine/pytorch_engine.py" line_range="21-30" />
<code_context>
+def _sample_next_token(logits, temperature, top_p):
</code_context>
<issue_to_address>
**suggestion:** Entropy for top-p sampling is computed on the truncated distribution, which may be misleading.
In the `top_p < 1.0` path, `entropy` is computed from the truncated, renormalized `filtered` distribution, while in the else path it’s computed from full `probs`. This means entropy depends on whether top-p is used, which may be surprising if callers expect it to reflect the model’s true next-token distribution. Please either document that entropy is over the truncated distribution or compute it from the original `probs` in both branches (e.g., before applying top-p).
Suggested implementation:
```python
def _sample_next_token(logits, temperature, top_p):
scaled = logits / max(float(temperature), 1e-7)
probs = F.softmax(scaled.float(), dim=-1)
# Entropy is computed on the full next-token distribution (before any top-p truncation)
log_probs = probs.clamp_min(1e-12).log()
entropy = -(probs * log_probs).sum(dim=-1)
if top_p < 1.0:
```
I can’t see the rest of `_sample_next_token`, but to fully implement the suggestion you should:
1. Remove any existing entropy computation that uses the truncated distribution, e.g. something like:
- `entropy = -(filtered * filtered.clamp_min(1e-12).log()).sum(dim=-1)` inside the `if top_p < 1.0:` block, or
- any entropy expression that uses `filtered` or renormalized top-p probabilities.
2. Ensure both the `top_p < 1.0` and `else` paths use the `entropy` defined from `probs` above (for returning, logging, or attaching to `SampleResult`).
3. If the function returns entropy, make sure the return signature and call sites are unchanged other than now using the shared `entropy` value.
4. Optionally, you may want to add or update the function (or docstring/comment) to document that `entropy` is computed on the full distribution before top-p sampling.
</issue_to_address>
### Comment 2
<location path="retrain/scaleway_backend.py" line_range="171-185" />
<code_context>
+ # Internals
+ # ------------------------------------------------------------------
+
+ def _wait_healthy(self) -> None:
+ deadline = time.monotonic() + self._health_timeout_s
+ for url in (f"{self._inference_url}/health", f"{self._training_url}/health"):
+ logger.info("Waiting for %s …", url)
+ while True:
+ try:
+ r = self._client.get(url, timeout=5)
+ if r.status_code == 200:
+ break
+ except Exception:
+ pass
+ if time.monotonic() > deadline:
+ raise RuntimeError(f"Timeout waiting for {url} to become healthy")
+ time.sleep(self._health_poll_s)
+ logger.info("Both services healthy.")
+
+ def _sample_one(
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Shared deadline for inference and training health checks can shorten the effective timeout for the second service.
Because the same `deadline` is shared across both health checks, the second service may get much less than `health_timeout_s` if the first one is slow to become healthy. This can cause premature failures for slower or flaky boots. Consider computing a fresh deadline per service (e.g., inside the loop) so each gets the full timeout budget.
```suggestion
def _wait_healthy(self) -> None:
for url in (f"{self._inference_url}/health", f"{self._training_url}/health"):
deadline = time.monotonic() + self._health_timeout_s
logger.info("Waiting for %s …", url)
while True:
try:
r = self._client.get(url, timeout=5)
if r.status_code == 200:
break
except Exception:
pass
if time.monotonic() > deadline:
raise RuntimeError(f"Timeout waiting for {url} to become healthy")
time.sleep(self._health_poll_s)
logger.info("Both services healthy.")
```
</issue_to_address>
### Comment 3
<location path="retrain/scaleway_backend.py" line_range="234-239" />
<code_context>
+ raise RuntimeError("Scaleway inference logprobs.content entry is missing logprob.")
+ logprobs.append(float(logprob))
+
+ if completion_text:
+ expected_ids = self._get_tokenizer().encode(completion_text, add_special_tokens=False)
+ if expected_ids and len(expected_ids) != len(token_ids):
+ raise RuntimeError(
+ "Scaleway inference returned logprobs.content length "
+ f"{len(token_ids)} but completion text encodes to {len(expected_ids)} tokens."
+ )
+ if len(token_ids) != len(logprobs):
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Strict length equality check between logprobs tokens and re-encoded completion text may be too brittle.
Raising `RuntimeError` on any mismatch between `expected_ids` and `token_ids` means minor tokenization or formatting differences (e.g., whitespace, provider changes) will abort otherwise usable outputs. Consider treating this as a soft failure (warning/metric) and falling back—either trusting the provider’s `logprobs` when lengths differ, or only re-encoding when lengths already match—so harmless upstream changes don’t break training runs.
Suggested implementation:
```python
if completion_text:
# Best-effort sanity check: only use re-encoding when it naturally matches the provider output.
# Length mismatches can happen due to tokenizer/provider changes; treat them as a soft failure
# and trust the provider's token_ids/logprobs instead of aborting.
expected_ids = self._get_tokenizer().encode(completion_text, add_special_tokens=False)
if expected_ids and len(expected_ids) != len(token_ids):
expected_ids = []
```
If you have a centralized logging/metrics setup, you may also want to:
1. Log a warning (or emit a metric) when `expected_ids` length differs from `token_ids` to monitor how often this soft failure happens.
2. That would involve importing your logger in this module and adding a log call inside the `if expected_ids and len(expected_ids) != len(token_ids):` block.
</issue_to_address>
### Comment 4
<location path="retrain/scaleway/terraform_runner.py" line_range="123-132" />
<code_context>
+ v("max_model_len", self._max_model_len),
+ ]
+
+ def _run_tf(self, *args: str) -> None:
+ cmd = ["terraform", *args]
+ env = {**os.environ, "TF_DATA_DIR": str(self._state_dir)}
+ result = subprocess.run( # nosec B603 — list argv, shell=False, no user input reaches here
+ cmd,
+ cwd=str(self._tf_dir),
+ env=env,
+ capture_output=False,
+ text=True,
+ )
+ if result.returncode != 0:
+ raise TerraformError(f"terraform {args[0]} exited with code {result.returncode}")
+
+ def _read_outputs(self) -> dict[str, object]:
</code_context>
<issue_to_address>
**suggestion:** Terraform error reporting omits stderr, making failures harder to diagnose.
Currently, failures only expose the exit code, while stderr is only printed to the console and not available to callers. For non-interactive or remote use, this makes diagnosing issues difficult. Please adjust the subprocess call so that, on failure, stderr is captured (even if truncated) and included in the `TerraformError` message to improve debuggability.
Suggested implementation:
```python
def _run_tf(self, *args: str) -> None:
cmd = ["terraform", *args]
env = {**os.environ, "TF_DATA_DIR": str(self._state_dir)}
result = subprocess.run( # nosec B603 — list argv, shell=False, no user input reaches here
cmd,
cwd=str(self._tf_dir),
env=env,
capture_output=True,
text=True,
)
if result.returncode != 0:
stderr = (result.stderr or "").strip()
# Limit stderr size to keep error messages manageable
max_stderr_chars = 2000
if len(stderr) > max_stderr_chars:
stderr = stderr[:max_stderr_chars] + "\n...[truncated]"
base_msg = f"terraform {args[0]} exited with code {result.returncode}"
if stderr:
raise TerraformError(f"{base_msg}; stderr:\n{stderr}")
raise TerraformError(base_msg)
```
You may also want to align `_read_outputs`'s error handling with `_run_tf` by including its `stderr` content in any raised `TerraformError`, using the same truncation pattern for consistency.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+21
to
+30
| def _sample_next_token(logits, temperature, top_p): | ||
| scaled = logits / max(float(temperature), 1e-7) | ||
| probs = F.softmax(scaled.float(), dim=-1) | ||
|
|
||
| if top_p < 1.0: | ||
| sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1) | ||
| cumulative = torch.cumsum(sorted_probs, dim=-1) | ||
| remove = cumulative - sorted_probs > top_p | ||
| filtered = sorted_probs.masked_fill(remove, 0.0) | ||
| filtered = filtered / filtered.sum(dim=-1, keepdim=True).clamp_min(1e-12) |
There was a problem hiding this comment.
suggestion: Entropy for top-p sampling is computed on the truncated distribution, which may be misleading.
In the top_p < 1.0 path, entropy is computed from the truncated, renormalized filtered distribution, while in the else path it’s computed from full probs. This means entropy depends on whether top-p is used, which may be surprising if callers expect it to reflect the model’s true next-token distribution. Please either document that entropy is over the truncated distribution or compute it from the original probs in both branches (e.g., before applying top-p).
Suggested implementation:
def _sample_next_token(logits, temperature, top_p):
scaled = logits / max(float(temperature), 1e-7)
probs = F.softmax(scaled.float(), dim=-1)
# Entropy is computed on the full next-token distribution (before any top-p truncation)
log_probs = probs.clamp_min(1e-12).log()
entropy = -(probs * log_probs).sum(dim=-1)
if top_p < 1.0:I can’t see the rest of _sample_next_token, but to fully implement the suggestion you should:
- Remove any existing entropy computation that uses the truncated distribution, e.g. something like:
entropy = -(filtered * filtered.clamp_min(1e-12).log()).sum(dim=-1)inside theif top_p < 1.0:block, or- any entropy expression that uses
filteredor renormalized top-p probabilities.
- Ensure both the
top_p < 1.0andelsepaths use theentropydefined fromprobsabove (for returning, logging, or attaching toSampleResult). - If the function returns entropy, make sure the return signature and call sites are unchanged other than now using the shared
entropyvalue. - Optionally, you may want to add or update the function (or docstring/comment) to document that
entropyis computed on the full distribution before top-p sampling.
Comment on lines
+171
to
+185
| def _wait_healthy(self) -> None: | ||
| deadline = time.monotonic() + self._health_timeout_s | ||
| for url in (f"{self._inference_url}/health", f"{self._training_url}/health"): | ||
| logger.info("Waiting for %s …", url) | ||
| while True: | ||
| try: | ||
| r = self._client.get(url, timeout=5) | ||
| if r.status_code == 200: | ||
| break | ||
| except Exception: | ||
| pass | ||
| if time.monotonic() > deadline: | ||
| raise RuntimeError(f"Timeout waiting for {url} to become healthy") | ||
| time.sleep(self._health_poll_s) | ||
| logger.info("Both services healthy.") |
There was a problem hiding this comment.
suggestion (bug_risk): Shared deadline for inference and training health checks can shorten the effective timeout for the second service.
Because the same deadline is shared across both health checks, the second service may get much less than health_timeout_s if the first one is slow to become healthy. This can cause premature failures for slower or flaky boots. Consider computing a fresh deadline per service (e.g., inside the loop) so each gets the full timeout budget.
Suggested change
| def _wait_healthy(self) -> None: | |
| deadline = time.monotonic() + self._health_timeout_s | |
| for url in (f"{self._inference_url}/health", f"{self._training_url}/health"): | |
| logger.info("Waiting for %s …", url) | |
| while True: | |
| try: | |
| r = self._client.get(url, timeout=5) | |
| if r.status_code == 200: | |
| break | |
| except Exception: | |
| pass | |
| if time.monotonic() > deadline: | |
| raise RuntimeError(f"Timeout waiting for {url} to become healthy") | |
| time.sleep(self._health_poll_s) | |
| logger.info("Both services healthy.") | |
| def _wait_healthy(self) -> None: | |
| for url in (f"{self._inference_url}/health", f"{self._training_url}/health"): | |
| deadline = time.monotonic() + self._health_timeout_s | |
| logger.info("Waiting for %s …", url) | |
| while True: | |
| try: | |
| r = self._client.get(url, timeout=5) | |
| if r.status_code == 200: | |
| break | |
| except Exception: | |
| pass | |
| if time.monotonic() > deadline: | |
| raise RuntimeError(f"Timeout waiting for {url} to become healthy") | |
| time.sleep(self._health_poll_s) | |
| logger.info("Both services healthy.") |
Comment on lines
+234
to
+239
| if completion_text: | ||
| expected_ids = self._get_tokenizer().encode(completion_text, add_special_tokens=False) | ||
| if expected_ids and len(expected_ids) != len(token_ids): | ||
| raise RuntimeError( | ||
| "Scaleway inference returned logprobs.content length " | ||
| f"{len(token_ids)} but completion text encodes to {len(expected_ids)} tokens." |
There was a problem hiding this comment.
suggestion (bug_risk): Strict length equality check between logprobs tokens and re-encoded completion text may be too brittle.
Raising RuntimeError on any mismatch between expected_ids and token_ids means minor tokenization or formatting differences (e.g., whitespace, provider changes) will abort otherwise usable outputs. Consider treating this as a soft failure (warning/metric) and falling back—either trusting the provider’s logprobs when lengths differ, or only re-encoding when lengths already match—so harmless upstream changes don’t break training runs.
Suggested implementation:
if completion_text:
# Best-effort sanity check: only use re-encoding when it naturally matches the provider output.
# Length mismatches can happen due to tokenizer/provider changes; treat them as a soft failure
# and trust the provider's token_ids/logprobs instead of aborting.
expected_ids = self._get_tokenizer().encode(completion_text, add_special_tokens=False)
if expected_ids and len(expected_ids) != len(token_ids):
expected_ids = []If you have a centralized logging/metrics setup, you may also want to:
- Log a warning (or emit a metric) when
expected_idslength differs fromtoken_idsto monitor how often this soft failure happens. - That would involve importing your logger in this module and adding a log call inside the
if expected_ids and len(expected_ids) != len(token_ids):block.
Comment on lines
+123
to
+132
| def _run_tf(self, *args: str) -> None: | ||
| cmd = ["terraform", *args] | ||
| env = {**os.environ, "TF_DATA_DIR": str(self._state_dir)} | ||
| result = subprocess.run( # nosec B603 — list argv, shell=False, no user input reaches here | ||
| cmd, | ||
| cwd=str(self._tf_dir), | ||
| env=env, | ||
| capture_output=False, | ||
| text=True, | ||
| ) |
There was a problem hiding this comment.
suggestion: Terraform error reporting omits stderr, making failures harder to diagnose.
Currently, failures only expose the exit code, while stderr is only printed to the console and not available to callers. For non-interactive or remote use, this makes diagnosing issues difficult. Please adjust the subprocess call so that, on failure, stderr is captured (even if truncated) and included in the TerraformError message to improve debuggability.
Suggested implementation:
def _run_tf(self, *args: str) -> None:
cmd = ["terraform", *args]
env = {**os.environ, "TF_DATA_DIR": str(self._state_dir)}
result = subprocess.run( # nosec B603 — list argv, shell=False, no user input reaches here
cmd,
cwd=str(self._tf_dir),
env=env,
capture_output=True,
text=True,
)
if result.returncode != 0:
stderr = (result.stderr or "").strip()
# Limit stderr size to keep error messages manageable
max_stderr_chars = 2000
if len(stderr) > max_stderr_chars:
stderr = stderr[:max_stderr_chars] + "\n...[truncated]"
base_msg = f"terraform {args[0]} exited with code {result.returncode}"
if stderr:
raise TerraformError(f"{base_msg}; stderr:\n{stderr}")
raise TerraformError(base_msg)You may also want to align _read_outputs's error handling with _run_tf by including its stderr content in any raised TerraformError, using the same truncation pattern for consistency.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds the retrain-side Gemma4 local-backend integration needed for Gemma 4 E4B QAT LoRA trainability on a single 4070 Ti.
This draft PR is opened from the existing
pr-3-scalewaybranch, so it also contains the prior Scaleway branch history already present before the Gemma integration commit.Validation
PYTHONPYCACHEPREFIX=/tmp/retrain-pycache python3 -m py_compile retrain/local_train_helper.py retrain/inference_engine/__init__.py retrain/inference_engine/pytorch_engine.pygit diff --checkLocalTrainHelper.sample_with_entropyandtrain_stepwith nonzero logprobs/entropy and shared train/infer modelSummary by Sourcery
Add Scaleway cloud backend and Gemma4 text-only local training/sampling support while improving local backend sharing and tooling.
New Features:
Enhancements:
Build:
Documentation:
Tests: