From ee675b488f9e2fbf0f787a5658482bd640648c02 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 30 Oct 2026 08:52:33 +0000 Subject: [PATCH] fix(cs): initialize score heads and correct training metric reporting --- cookbook/client/twinkle/prepare_seq_cls_model.py | 7 ++++--- cookbook/client/twinkle/reranker.py | 5 +++-- cookbook/client/twinkle/seq_cls.py | 5 +++-- src/twinkle/server/model/app.py | 4 ++++ tests/server/integration/task_http_validation.py | 7 ++++--- tests/server/model/test_task_head.py | 10 ++++++++++ 6 files changed, 28 insertions(+), 10 deletions(-) diff --git a/cookbook/client/twinkle/prepare_seq_cls_model.py b/cookbook/client/twinkle/prepare_seq_cls_model.py index 7844c21b..1f80dc96 100644 --- a/cookbook/client/twinkle/prepare_seq_cls_model.py +++ b/cookbook/client/twinkle/prepare_seq_cls_model.py @@ -35,10 +35,11 @@ # The classifier scores the last non-padding token; batched inputs need the pad id to find it. text_config.pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id -torch.manual_seed(42) -# The load report lists the vision weights as UNEXPECTED and score.weight as MISSING: -# both are expected, the vision tower is dropped and the new head is initialized here. +# The load report lists the vision weights as UNEXPECTED and score.weight as MISSING. model = AutoModelForSequenceClassification.from_pretrained(model_dir, config=text_config, torch_dtype='auto') +# Initialize the new head explicitly instead of relying on implicit missing-weight initialization. +torch.manual_seed(42) +torch.nn.init.normal_(model.score.weight, mean=0.0, std=text_config.initializer_range) model.save_pretrained(args.output_dir) tokenizer.save_pretrained(args.output_dir) print(f'Saved {type(model).__name__} with {args.num_labels} label(s) to {args.output_dir}') # noqa: T201 diff --git a/cookbook/client/twinkle/reranker.py b/cookbook/client/twinkle/reranker.py index d7fa7098..3af50013 100644 --- a/cookbook/client/twinkle/reranker.py +++ b/cookbook/client/twinkle/reranker.py @@ -84,9 +84,10 @@ def train(): # Step 5: training loop. task='reranker' returns one score per pair. for step in range(MAX_STEPS): - output = model.forward_backward(inputs=batch, task='reranker') + model.forward_backward(inputs=batch, task='reranker') model.clip_grad_and_step(max_grad_norm=1.0) - logger.info(f'Step {step}: loss={output.result["loss"]:.4f}') + metric = model.calculate_metric(is_training=True) + logger.info(f'Step {step}: {metric.result}') # Step 6: relevant documents should now score higher, then save the adapter. scores = model.forward_only(inputs=batch, task='reranker', return_logits=True).result['logits'] diff --git a/cookbook/client/twinkle/seq_cls.py b/cookbook/client/twinkle/seq_cls.py index c01f1744..84b516ba 100644 --- a/cookbook/client/twinkle/seq_cls.py +++ b/cookbook/client/twinkle/seq_cls.py @@ -90,9 +90,10 @@ def train(): # Step 5: training loop. task='seq_cls' returns one [num_labels] score row per text. for step in range(MAX_STEPS): - output = model.forward_backward(inputs=batch, task='seq_cls') + model.forward_backward(inputs=batch, task='seq_cls') model.clip_grad_and_step(max_grad_norm=1.0) - logger.info(f'Step {step}: loss={output.result["loss"]:.4f}') + metric = model.calculate_metric(is_training=True) + logger.info(f'Step {step}: {metric.result}') # Step 6: check the predictions, then save the adapter and optimizer state. logits = model.forward_only(inputs=batch, task='seq_cls', return_logits=True).result['logits'] diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/model/app.py index 908dddbd..80c4d50e 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/model/app.py @@ -101,6 +101,10 @@ def _make_model_class(base: type) -> type: def _make_transformers_model(kw: dict[str, Any]) -> Any: train_mode = kw.pop('train_mode', 'lora') + # The head comes from the persisted config checked by ``_validate_model_task``; HF model + # constructors do not accept these deployment keys, unlike Megatron's bridge config. + kw.pop('task_type', None) + kw.pop('num_labels', None) # ``model_cls`` arrives as plain data and is resolved here, so the class object never reaches the # constructor -- it would either be rejected by the remote-class layer in safe mode # (``TWINKLE_TRUST_REMOTE_CODE=0``) or leak into every worker's constructor kwargs. The schema only diff --git a/tests/server/integration/task_http_validation.py b/tests/server/integration/task_http_validation.py index 7a051db2..3f4468b5 100644 --- a/tests/server/integration/task_http_validation.py +++ b/tests/server/integration/task_http_validation.py @@ -70,11 +70,12 @@ def step(model): kwargs = {'task': args.task, 'gradient_accumulation_steps': 2} if args.backend == 'megatron': kwargs['micro_batch_size'] = 6 - result = model.forward_backward(inputs=copy.deepcopy(batch), **kwargs).result - assert np.isfinite(result['loss']) + model.forward_backward(inputs=copy.deepcopy(batch), **kwargs) # Call after every backward; the backend gates updates at its GA boundary. model.clip_grad_and_step() - return float(result['loss']) + loss = float(model.calculate_metric(is_training=True).result['loss']) + assert np.isfinite(loss) + return loss def save(model, name, samples): return model.save(name, save_optimizer=True, consumed_train_samples=samples) diff --git a/tests/server/model/test_task_head.py b/tests/server/model/test_task_head.py index 681d6e52..055b2ee9 100644 --- a/tests/server/model/test_task_head.py +++ b/tests/server/model/test_task_head.py @@ -99,3 +99,13 @@ def tokenizer(text, **kwargs): parsed = TypeAdapter(WireInputBatch).validate_python(batch) assert len(parsed) == 12 assert all(type(label) is int for row in parsed for label in row.labels) + + +def test_transformers_backend_does_not_forward_head_keys(monkeypatch): + from twinkle.server.model import app + from twinkle.server.model.backends import transformers_model + + received = {} + monkeypatch.setattr(transformers_model, 'TwinkleCompatTransformersModel', lambda **kwargs: received.update(kwargs)) + app._make_transformers_model({'model_id': 'unused', 'task_type': 'seq_cls', 'num_labels': 3}) + assert received == {'model_id': 'unused'}