convaiinnovations

/

laya

Like 5.78k

Versions

• laya 0.4.0 (pip), torch 2.14.0+cu130, transformers 5.17.0, Python 3.12, Ubuntu, RTX 3090 Ti
• Base checkpoint: convaiinnovations/laya @ 7b928d82, subfolder typed-decisions
• The same data and arguments trained without error on laya 0.3.28.

Command

laya-train --data dev.jsonl --base <snapshot>/typed-decisions --out ckpt --device cuda --loss soft-ce --shuffle-options --seed 0

Error

laya-train: unexpected error: Expected all tensors to be on the same device, but got index is on cuda:0,
different from other tensors on cpu (when checking argument in method wrapper_CUDA__index_select)

Traceback (from calling laya.train.finetune directly, because the CLI hides it)

laya/train.py, line 1086, in finetune
    base_records = calibration_records(model, tok, eval_items, dev, max_len, head_max_len, parallel=parallel)
laya/train.py, line 718, in calibration_records
    logits = _forward(model, batch, device, amp=False, detach_encoder=False).cpu().numpy()
laya/train.py, line 617, in _forward
    logits, _act = model(*args, **kwargs)
laya/common.py, line 612, in forward
    h = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
... modeling_modernbert.py, line 70: self.tok_embeddings(input_ids)
RuntimeError: Expected all tensors to be on the same device ...

Cause
load_checkpoint() returns the model on CPU (its docstring says so). The model is moved to the device only inside train_model() (model.to(device).train(), train.py line 638). The pre-training evaluation of the base checkpoint added in 0.4.0 (train.py around line 1086) runs before that, sending CUDA inputs to a CPU model. Any non-CPU device fails; --device cpu would not.

One-line fix (verified locally)
In finetune(), move the model before the base evaluation:

        model.to(dev)
        base_records = calibration_records(model, tok, eval_items, dev, max_len, head_max_len, parallel=parallel)

With this line, training proceeds past that point on CUDA. train_model() moving the model again is harmless.

Also seen (minor, may be expected): loading the typed-decisions checkpoint warns this checkpoint ships invalid temperatures or values outside [0.5, 5]; using choice:11+=0.1005... -> 0.5.

Thanks for Laya, I love it and I've done some cool things with it.