es-ude / es-ude/OnDeviceTraining
training_loop: loss-grad seed wire has no storage-quantization knob — inherits model-output dtype
- Dominant language
- C
- Stars
- 1
- Forks
- 3
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 8
Description
## Problem
The topmost backward wire — the tensor the loss backward writes into, feeding the first backward layer — is the only wire in the training loop without a user-facing storage-quantization knob.
`calculateGradsImpl` allocates it with a NULL wire config (src/userApi/training_loop/calculate_grads/CalculateGradsSequential.c:75):
```c
initGradTensor(&gradNext, layerOutputs[modelSize], NULL);
lossFns.backward(layerOutputs[modelSize], label, &gradNext);
```
and `initGradTensor` falls back to the model output's quantization (CalculateGradsSequential.c:244), cloning its full config including `qMaxBits`/`roundingMode` (:266-279). Every layer-produced dx wire instead gets the producer's declared `propLossQ` via `backwardWireQ(model[i])` (CalculateGradsSequential.c:89, src/userApi/LayerConfigAccess.c:63). `lossConfig_t` has no quantization field (src/loss_functions/include/LossFunction.h:10-14) — the gap is even documented as such in the `backwardWireQ` doc comment (LayerConfigAccess.c:59-62).
Concrete consequences:
- **Head-grad storage is coupled to output storage.** An MSE head with a SYM_INT32 `outputQ` silently gets a SYM_INT32 seed grad at the *output's* width/rounding — a config chosen for activations, not gradients — and the user cannot pick FLOAT32 (or a different width/rounding) for the seed without changing the model-output storage itself. The reverse also holds: with a FLOAT32 output, a fully-quantized-wires configuration (every layer's `outputQ`/`propLossQ` SYM) stays pinned FLOAT32 at the head, silently.
- **In CE mode the seed is doubly unconfigurable.** The fused CE+Softmax shortcut skips the softmax backward (CalculateGradsSequential.c:70-72), so the seed *replaces* softmax's dx wire and softmax's own `propLossQ` knob is bypassed — there is no quantization control point at all between the loss and the first real backward layer. (The fused CE backward's missing SYM compute arm is a separate concern, owned by #206; this issue is only about the seed-wire *allocation* knob.)
## Proposal
Add an optional seed-wire storage config, `quantization_t *lossGradQ`, to `lossConfig_t` (next to `backwardReduction`/`classWeights`), defaulting to NULL = current behavior (inherit model-output dtype). Wiring is one line: pass it instead of NULL at CalculateGradsSequential.c:75 — `initGradTensor` already accepts a wireQ and the loss backwards already dispatch/convert on tensor dtype (MSE's SYM arm writes through `convertTensor`). Keep scope to the dtypes `initGradTensor` handles today (FLOAT32, SYM_INT32; fail fast otherwise, matching the existing switch), and update the LayerConfigAccess.c:59-62 comment when the knob ships.
Refs: #300 (axis 6 — the wire-storage memo names this boundary; the epic scopes fixing it out to a dedicated issue), #206 (adjacent: CE loss-head SYM compute arms)
Contributor guide
Research direction
Start at calculateGradsImpl in src/userApi/training_loop/calculate_grads/CalculateGradsSequential.c, then read initGradTensor there and the lossConfig_t definition in src/loss_functions/include/LossFunction.h. Compare the seed allocation with backwardWireQ in src/userApi/LayerConfigAccess.c. Done means an optional lossGradQ controls seed-wire storage while NULL preserves current inheritance, with the related comment updated and unsupported dtypes handled consistently.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- c
- Domain
- embedded-iot
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 74/100