Repository navigation
Expand file tree
/
Copy pathtrainer.py
More file actions
2070 lines (1847 loc) · 85 KB
/
Copy pathtrainer.py
File metadata and controls
2070 lines (1847 loc) · 85 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
"""
GLiNER2 World-Class Trainer
===========================
Production-grade training infrastructure with flexible data input.
Supported Data Formats:
-----------------------
1. Single JSONL file path (str or Path)
2. List of JSONL file paths
3. List of InputExample objects
4. TrainingDataset object
5. List of raw dict records ({"input": ..., "output": ...} format)
Basic Examples:
--------------
>>> from gliner2.training.data import InputExample, TrainingDataset
>>> from gliner2.training.trainer import TrainingConfig, GLiNER2Trainer
>>>
>>> # 1. From list of InputExample
>>> examples = [
... InputExample(text="John works at Google.", entities={"person": ["John"], "company": ["Google"]}),
... InputExample(text="Apple released iPhone.", entities={"company": ["Apple"], "product": ["iPhone"]}),
... ]
>>> trainer = GLiNER2Trainer(model, config)
>>> trainer.train(train_data=examples)
>>>
>>> # 2. From JSONL file(s)
>>> trainer.train(train_data="train.jsonl")
>>> trainer.train(train_data=["train1.jsonl", "train2.jsonl"])
>>>
>>> # 3. From TrainingDataset
>>> dataset = TrainingDataset.load("train.jsonl")
>>> trainer.train(train_data=dataset)
"""
from __future__ import annotations
import contextlib
import gc
import hashlib
import json
import logging
import math
import os
import random
import shutil
import sys
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field, fields as dataclass_fields, is_dataclass, asdict
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import numpy as np
import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader, Dataset, DistributedSampler
import torch.distributed as dist
from tqdm.auto import tqdm
from gliner2.processor import SchemaTransformer, SamplingConfig
from gliner2.utils.device_errors import is_fatal_device_error
from gliner2.utils.sync_probe import count_cuda_syncs
# Import training data classes
from gliner2.training.config import TrainingConfig
from gliner2.training.data import (
InputExample, TrainingDataset, DataValidationError,
DataFormat, detect_data_format, DataLoader_Factory, TrainDataInput
)
from gliner2.training.sampler import (
DistributedLengthGroupedSampler,
LengthGroupedSampler,
)
logger = logging.getLogger(__name__)
# =============================================================================
# Configuration
# =============================================================================
# =============================================================================
# Dataset
# =============================================================================
class ExtractorDataset(Dataset):
"""
Dataset for GLiNER2 training with multi-format support.
Supports all formats through DataLoader_Factory:
- JSONL file path(s)
- List of InputExample objects
- TrainingDataset object
- List of raw dict records
Examples
--------
>>> # From JSONL
>>> dataset = ExtractorDataset("train.jsonl")
>>> # From multiple JSONL files
>>> dataset = ExtractorDataset(["train1.jsonl", "train2.jsonl"])
>>> # From InputExample list
>>> dataset = ExtractorDataset(examples)
"""
def __init__(
self,
data: TrainDataInput,
max_samples: int = -1,
shuffle: bool = True,
seed: int = 42,
validate: bool = False,
):
"""
Initialize dataset from various input formats.
Parameters
----------
data : TrainDataInput
Training data in any supported format.
max_samples : int, default=-1
Maximum samples to use (-1 = all).
shuffle : bool, default=True
Whether to shuffle the data.
seed : int, default=42
Random seed for shuffling.
validate : bool, default=False
Whether to validate the data. Validation is always strict:
checks that entity spans, relation values, and structure
field values exist in the text.
"""
self.data = DataLoader_Factory.load(
data=data,
max_samples=max_samples,
shuffle=shuffle,
seed=seed,
validate=validate,
)
self.lengths = self._load_or_compute_lengths(data)
@staticmethod
def _record_text(record: Dict[str, Any]) -> str:
return str(record.get("input", record.get("text", "")))
def _length_fingerprint(self) -> str:
digest = hashlib.sha256()
for record in self.data:
text = self._record_text(record)
digest.update(len(text).to_bytes(8, "little"))
digest.update(text.encode("utf-8"))
return digest.hexdigest()
def _load_or_compute_lengths(self, source: TrainDataInput) -> Tuple[int, ...]:
"""Return cached inexpensive token-count estimates for bucketing.
A single JSONL source gets a ``.lengths.npy`` sidecar plus fingerprint
metadata. In-memory and multi-source datasets are cached only for the
lifetime of this object. Sidecar failures are deliberately non-fatal.
"""
fingerprint = self._length_fingerprint()
source_path = (
Path(source).expanduser()
if isinstance(source, (str, Path))
else None
)
cache_path = (
source_path.with_suffix(source_path.suffix + ".lengths.npy")
if source_path is not None else None
)
fingerprint_path = (
source_path.with_suffix(source_path.suffix + ".lengths.json")
if source_path is not None else None
)
if cache_path is not None and fingerprint_path is not None:
try:
metadata = json.loads(fingerprint_path.read_text())
cached = np.load(cache_path, allow_pickle=False)
if (
metadata.get("fingerprint") == fingerprint
and cached.ndim == 1
and cached.shape[0] == len(self.data)
):
return tuple(int(value) for value in cached)
except (OSError, ValueError, json.JSONDecodeError):
pass
lengths = tuple(
max(1, len(self._record_text(record).split()))
for record in self.data
)
if cache_path is not None and fingerprint_path is not None:
try:
np.save(cache_path, np.asarray(lengths, dtype=np.int32))
fingerprint_path.write_text(json.dumps({"fingerprint": fingerprint}))
except OSError:
logger.debug("Could not persist dataset length cache", exc_info=True)
return lengths
def __len__(self) -> int:
return len(self.data)
def __getitem__(self, idx: int) -> Tuple[str, Dict]:
record = self.data[idx]
# Handle both formats
if "input" in record:
return record["input"], record["output"]
else:
return record["text"], record["schema"]
# Factory methods for explicit creation
@classmethod
def from_jsonl(cls, paths: Union[str, Path, List], **kwargs) -> 'ExtractorDataset':
"""Create from JSONL file(s)."""
return cls(paths, **kwargs)
@classmethod
def from_examples(cls, examples: List[InputExample], **kwargs) -> 'ExtractorDataset':
"""Create from list of InputExample."""
return cls(examples, **kwargs)
@classmethod
def from_training_dataset(cls, dataset: TrainingDataset, **kwargs) -> 'ExtractorDataset':
"""Create from TrainingDataset."""
return cls(dataset, **kwargs)
@classmethod
def from_dicts(cls, dicts: List[Dict], **kwargs) -> 'ExtractorDataset':
"""Create from list of dicts."""
return cls(dicts, **kwargs)
# =============================================================================
# Collator
# =============================================================================
class ExtractorCollator:
"""Data collator that converts raw records to model inputs."""
def __init__(
self, processor: SchemaTransformer, is_training: bool = True,
max_len=None, architecture: str = "span",
max_gold_per_query: Optional[int] = 32,
build_targets: Optional[bool] = None,
on_capacity_exceeded: str = "raise",
allow_invalid_samples: bool = False,
dedupe_contained_entity_spans: bool = True,
):
self.processor = processor
self.is_training = is_training
self.max_len = max_len
self.architecture = architecture
self.max_gold_per_query = max_gold_per_query
# For an eval collator (``is_training=False``) set ``build_targets=True``
# so a supervised eval loss can be computed while the model runs in eval
# mode. Defaults to ``is_training`` (plain inference builds no targets).
self.build_targets = build_targets
# Gold-capacity overflow policy (raise | truncate_with_warning |
# skip_sample); defaults to the no-silent-loss "raise".
self.on_capacity_exceeded = on_capacity_exceeded
self.allow_invalid_samples = allow_invalid_samples
# Filter phantom contained-span entity matches (see
# gliner2.processing.boundary_preprocessing._dedupe_contained_spans);
# defaults to the corrected behavior.
self.dedupe_contained_entity_spans = dedupe_contained_entity_spans
def __call__(self, batch: List[Tuple[str, Dict]]):
"""
Convert batch of (text, schema) tuples to PreprocessedBatch.
Args:
batch: List of (text, schema) tuples from dataset
Returns:
PreprocessedBatch ready for model.forward()
"""
if self.is_training:
kwargs = dict(
max_len=self.max_len, architecture=self.architecture,
max_gold_per_query=self.max_gold_per_query,
on_capacity_exceeded=self.on_capacity_exceeded,
dedupe_contained_entity_spans=self.dedupe_contained_entity_spans,
)
if self.allow_invalid_samples:
kwargs.update(
error_policy="skip",
ignore_missing_entities=True,
)
return self.processor.collate_fn_train(batch, **kwargs)
else:
kwargs = dict(
max_len=self.max_len, architecture=self.architecture,
build_targets=self.build_targets,
max_gold_per_query=self.max_gold_per_query,
on_capacity_exceeded=self.on_capacity_exceeded,
dedupe_contained_entity_spans=self.dedupe_contained_entity_spans,
)
if self.allow_invalid_samples:
kwargs.update(
error_policy="skip",
ignore_missing_entities=True,
)
return self.processor.collate_fn_inference(batch, **kwargs)
# =============================================================================
# Metrics
# =============================================================================
@dataclass
class TrainingMetrics:
"""Container for training metrics."""
loss: float = 0.0
classification_loss: float = 0.0
structure_loss: float = 0.0
count_loss: float = 0.0
learning_rate: float = 0.0
epoch: float = 0.0
step: int = 0
samples_seen: int = 0
throughput: float = 0.0
def to_dict(self) -> Dict[str, float]:
return asdict(self)
# =============================================================================
# Scheduler Factory
# =============================================================================
def get_scheduler(optimizer, scheduler_type, num_training_steps, num_warmup_steps, num_cycles=0.5):
"""Create learning rate scheduler."""
def lr_lambda_linear(step):
if step < num_warmup_steps:
return float(step) / float(max(1, num_warmup_steps))
return max(0.0, float(num_training_steps - step) / float(max(1, num_training_steps - num_warmup_steps)))
def lr_lambda_cosine(step):
if step < num_warmup_steps:
return float(step) / float(max(1, num_warmup_steps))
progress = float(step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
def lr_lambda_cosine_restarts(step):
if step < num_warmup_steps:
return float(step) / float(max(1, num_warmup_steps))
progress = float(step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((num_cycles * progress) % 1.0))))
def lr_lambda_constant(step):
if step < num_warmup_steps:
return float(step) / float(max(1, num_warmup_steps))
return 1.0
schedulers = {
"linear": lr_lambda_linear,
"cosine": lr_lambda_cosine,
"cosine_restarts": lr_lambda_cosine_restarts,
"constant": lr_lambda_constant,
}
if scheduler_type not in schedulers:
raise ValueError(f"Unknown scheduler: {scheduler_type}")
return LambdaLR(optimizer, schedulers[scheduler_type])
# =============================================================================
# Main Trainer
# =============================================================================
class ExtractorTrainer:
"""
World-class trainer for GLiNER2 with flexible multi-format data input.
Architecture-neutral: drives both the span and boundary architectures. The
legacy name ``GLiNER2Trainer`` is preserved as an alias.
Parameters
----------
model : nn.Module
The GLiNER2 model to train.
config : TrainingConfig
Training configuration.
processor : SchemaTransformer, optional
Schema processor. If None, uses model.processor.
train_data : TrainDataInput, optional
Training data (can be provided here or in train()).
eval_data : TrainDataInput, optional
Evaluation data.
compute_metrics : Callable, optional
Custom metrics function.
Supported Data Formats
----------------------
- Single JSONL file path (str or Path)
- List of JSONL file paths
- List of InputExample objects
- TrainingDataset object
- List of raw dict records
Examples
--------
>>> # With InputExample list
>>> examples = [InputExample(...), InputExample(...)]
>>> trainer = GLiNER2Trainer(model, config)
>>> trainer.train(train_data=examples)
>>> # With JSONL file
>>> trainer.train(train_data="train.jsonl")
>>> # With multiple JSONL files
>>> trainer.train(train_data=["train1.jsonl", "train2.jsonl"])
>>> # With TrainingDataset
>>> dataset = TrainingDataset.load("train.jsonl")
>>> trainer.train(train_data=dataset)
"""
def __init__(
self,
model: nn.Module,
config: TrainingConfig,
processor: SchemaTransformer = None,
train_data: TrainDataInput = None,
eval_data: TrainDataInput = None,
compute_metrics: Optional[Callable] = None,
):
self.model = model
self.config = config
if config.bf16 and not torch.cuda.is_bf16_supported():
raise RuntimeError("bf16 was requested but this CUDA device does not support it")
if (
getattr(model, "architecture", "span") == "boundary"
and not getattr(config, "_precision_explicit", True)
):
config.fp16 = False
config.bf16 = True
logger.info("Boundary architecture defaulting to bf16 (fp16 disabled)")
if torch.cuda.is_available() and not torch.cuda.is_bf16_supported():
raise RuntimeError(
"The boundary architecture defaults to bf16, but this CUDA "
"device does not support it; choose fp32 explicitly."
)
torch.set_float32_matmul_precision(config.float32_matmul_precision)
if torch.cuda.is_available():
torch.backends.cuda.matmul.allow_tf32 = config.allow_tf32
torch.backends.cudnn.allow_tf32 = config.allow_tf32
self.processor = processor or getattr(model, 'processor', None)
if self.processor is None:
raise ValueError("Processor must be provided or model must have .processor attribute")
self.train_data = train_data
self.eval_data = eval_data
self.compute_metrics = compute_metrics
self._setup_seed()
self._setup_device()
self._setup_output_dir()
self._setup_logging()
self.global_step = 0
self.epoch = 0
self.best_metric = float('inf') if not config.greater_is_better else float('-inf')
self.patience_counter = 0
self.train_metrics_history = []
self.eval_metrics_history = []
self.optimizer = None
self.scheduler = None
self.scaler = None
self.wandb_run = None
self.progress_bar = None
# Most recent successful micro-batch outputs / grad norm, for logging.
self._last_train_outputs = None
self._last_grad_norm = None
self._skip_counter = None
self._loss_accum = None
self._loss_finite_flag = None
self._finite_grad_hook_handles = []
# LoRA state
self.lora_layers = {}
self._setup_lora()
base_model = self.model
if config.gradient_checkpointing:
encoder = getattr(base_model, "encoder", None)
enable = getattr(encoder, "gradient_checkpointing_enable", None)
if enable is None:
raise ValueError("encoder does not support gradient checkpointing")
# Non-reentrant checkpointing keeps adapter gradients when frozen
# embeddings produce inputs without requires_grad.
enable(gradient_checkpointing_kwargs={"use_reentrant": False})
if config.compile_model:
compile_method = getattr(base_model, "compile", None)
if compile_method is None:
raise ValueError("model does not support compile_model=True")
compile_method(dynamic=True)
self._install_finite_grad_hooks()
self._setup_distributed()
def _install_finite_grad_hooks(self) -> None:
"""Zero gradients from a device-detected non-finite loss before DDP."""
def sanitize(gradient):
flag = self._loss_finite_flag
if flag is None:
return gradient
return torch.where(flag, gradient, torch.zeros_like(gradient))
self._finite_grad_hook_handles = [
parameter.register_hook(sanitize)
for parameter in getattr(self.model, "parameters", lambda: ())()
if parameter.requires_grad
]
def _setup_seed(self):
seed = self.config.seed
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.use_deterministic_algorithms(self.config.deterministic)
if self.config.deterministic:
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
else:
torch.backends.cudnn.benchmark = True
def _setup_device(self):
if self.config.local_rank >= 0:
torch.cuda.set_device(self.config.local_rank)
self.device = torch.device("cuda", self.config.local_rank)
self.is_distributed = True
if not dist.is_initialized():
dist.init_process_group(backend="nccl", init_method="env://")
logger.info(f"Initialized distributed training: rank {dist.get_rank()}/{dist.get_world_size()}")
elif torch.cuda.is_available():
self.device = torch.device("cuda")
self.is_distributed = False
else:
self.device = torch.device("cpu")
self.is_distributed = False
if self.config.fp16 or self.config.bf16:
logger.warning("Mixed precision disabled on CPU")
self.config.fp16 = False
self.config.bf16 = False
self.model.to(self.device)
if not self.config.fp16 and not self.config.bf16:
# Some encoder backends preserve a reduced parameter dtype when
# loaded from a checkpoint. Full-precision training must normalize
# every module before its activations reach FP32 task heads.
self.model.float()
logger.info(f"Using device: {self.device}")
def _setup_output_dir(self):
self.output_dir = Path(self.config.output_dir)
self.logs_dir = self.output_dir / "logs"
if self.is_main_process:
self.output_dir.mkdir(parents=True, exist_ok=True)
self.logs_dir.mkdir(exist_ok=True)
self.config.save(str(self.output_dir / "training_config.json"))
def _setup_logging(self):
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
level=logging.INFO if self.is_main_process else logging.WARNING,
)
# W&B setup (HuggingFace style)
self.wandb_run = None
if self.config.report_to_wandb and self.is_main_process:
try:
import wandb
self.wandb_run = wandb.init(
project=self.config.wandb_project or self.config.experiment_name,
entity=self.config.wandb_entity,
name=self.config.wandb_run_name or f"{self.config.experiment_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
config=asdict(self.config),
tags=self.config.wandb_tags,
notes=self.config.wandb_notes,
dir=str(self.output_dir),
)
logger.info(f"W&B run: {self.wandb_run.url}")
except ImportError:
logger.warning("wandb not installed. Run: pip install wandb")
self.config.report_to_wandb = False
def _setup_lora(self):
"""Setup LoRA for parameter-efficient fine-tuning if enabled."""
if not self.config.use_lora:
logger.info("LoRA is disabled")
return
from gliner2.training.lora import PeftModel, _PeftLoraLayer
if isinstance(self.model, PeftModel):
from gliner2.training.session_adapters import validate_loaded_adapter
validate_loaded_adapter(self.model, self.config)
self.lora_layers = {n: m for n, m in self.model.named_modules() if isinstance(m, _PeftLoraLayer)}
self.model._lora_layers = self.lora_layers
return
for p in self.model.parameters():
p.requires_grad = False
logger.info("Froze all model parameters for LoRA training")
self.model = self.model.apply_lora(
r=self.config.lora_r, alpha=self.config.lora_alpha,
dropout=self.config.lora_dropout,
targets=self.config.lora_target_modules,
use_dora=self.config.lora_use_dora,
)
self.lora_layers = {n: m for n, m in self.model.named_modules() if isinstance(m, _PeftLoraLayer)}
self.model._lora_layers = self.lora_layers
lora_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in self.model.parameters())
pct = (lora_params / total_params * 100) if total_params > 0 else 0.0
logger.info(f"LoRA setup complete: {lora_params:,} trainable / {total_params:,} total ({pct:.2f}%)")
def _setup_distributed(self):
"""Setup distributed training if enabled."""
if self.is_distributed:
self.model = torch.nn.parallel.DistributedDataParallel(
self.model,
device_ids=[self.config.local_rank],
output_device=self.config.local_rank,
find_unused_parameters=self.config.ddp_find_unused_parameters,
static_graph=self.config.ddp_static_graph,
gradient_as_bucket_view=True,
broadcast_buffers=False,
)
logger.info("Wrapped model in DistributedDataParallel")
def _cleanup_distributed(self):
if self.is_distributed and dist.is_initialized():
dist.destroy_process_group()
if self.is_distributed and hasattr(self.model, "module"):
self.model = self.model.module
self.is_distributed = False
@property
def is_main_process(self) -> bool:
return not self.is_distributed or dist.get_rank() == 0
@staticmethod
def _safe_divide(numerator: float, denominator: float, default: float = 0.0) -> float:
"""Safely divide two numbers, returning default if denominator is zero."""
if denominator == 0:
return default
return numerator / denominator
def _record_failed_batch(
self,
batch: Any,
*,
data_loader_step: int,
error: BaseException,
) -> None:
"""Persist CPU-side sample contents without touching a failed device."""
try:
texts = list(getattr(batch, "original_texts", ()) or ())
schemas = list(getattr(batch, "original_schemas", ()) or ())
samples = []
for sample_index in range(max(len(texts), len(schemas))):
text = texts[sample_index] if sample_index < len(texts) else None
schema = schemas[sample_index] if sample_index < len(schemas) else None
canonical = json.dumps(
{"text": text, "schema": schema},
ensure_ascii=False,
sort_keys=True,
default=repr,
)
samples.append({
"sample_index": sample_index,
"sha256": hashlib.sha256(canonical.encode("utf-8")).hexdigest(),
"text": text,
"schema": schema,
})
rank = dist.get_rank() if self.is_distributed else 0
record = {
"timestamp": datetime.now().isoformat(),
"rank": rank,
"epoch": self.epoch,
"global_step": self.global_step,
"data_loader_step": data_loader_step,
"error_type": type(error).__name__,
"error": str(error),
"samples": samples,
}
path = self.output_dir / (
"failed_batches.jsonl"
if not self.is_distributed
else f"failed_batches.rank-{rank}.jsonl"
)
with path.open("a", encoding="utf-8") as handle:
handle.write(json.dumps(record, ensure_ascii=False, default=repr))
handle.write("\n")
handle.flush()
os.fsync(handle.fileno())
logger.error(
"Recorded failing batch with %d sample(s) in %s",
len(samples),
path,
)
except Exception:
# Diagnostics must never replace the original training exception.
logger.exception("Could not record failing batch diagnostics")
@classmethod
def _diagnostic_json_value(cls, value: Any) -> Any:
"""Convert nested batch state, including tensors, to lossless JSON data."""
if isinstance(value, torch.Tensor):
tensor = value.detach().cpu()
return {
"__type__": "tensor",
"dtype": str(tensor.dtype),
"shape": list(tensor.shape),
"values": tensor.tolist(),
}
if isinstance(value, np.ndarray):
return {
"__type__": "ndarray",
"dtype": str(value.dtype),
"shape": list(value.shape),
"values": value.tolist(),
}
if isinstance(value, np.generic):
return value.item()
if is_dataclass(value) and not isinstance(value, type):
result = {"__type__": type(value).__name__}
result.update({
item.name: cls._diagnostic_json_value(getattr(value, item.name))
for item in dataclass_fields(value)
})
return result
if isinstance(value, dict):
return {
str(key): cls._diagnostic_json_value(item)
for key, item in value.items()
}
if isinstance(value, (list, tuple)):
return [cls._diagnostic_json_value(item) for item in value]
if value is None or isinstance(value, (bool, int, float, str)):
return value
if hasattr(value, "__dict__"):
result = {"__type__": type(value).__name__}
result.update({
str(key): cls._diagnostic_json_value(item)
for key, item in vars(value).items()
})
return result
return repr(value)
def _record_debug_batch(
self,
batch: Any,
*,
data_loader_step: int,
micro_batch_in_window: int,
is_last_micro: bool,
) -> None:
"""Dump a complete pre-GPU batch for a configured diagnostic step."""
try:
rank = dist.get_rank() if self.is_distributed else 0
record = {
"timestamp": datetime.now().isoformat(),
"rank": rank,
"epoch": self.epoch,
"global_step_before_batch": self.global_step,
"data_loader_step": data_loader_step,
"micro_batch_in_window": micro_batch_in_window,
"is_last_micro": is_last_micro,
"batch": self._diagnostic_json_value(batch),
}
path = self.output_dir / (
"debug_batches.jsonl"
if not self.is_distributed
else f"debug_batches.rank-{rank}.jsonl"
)
with path.open("a", encoding="utf-8") as handle:
handle.write(json.dumps(record, ensure_ascii=False))
handle.write("\n")
handle.flush()
os.fsync(handle.fileno())
logger.warning(
"Recorded full pre-GPU batch for global step %d in %s",
self.global_step,
path,
)
except Exception:
# Diagnostics must never introduce a new training failure.
logger.exception("Could not record configured batch diagnostics")
def _get_model_config(self) -> Any:
"""Return the underlying model config, handling DDP-wrapped models."""
base_model = self.model.module if self.is_distributed and hasattr(self.model, "module") else self.model
return getattr(base_model, "config", None)
def _validate_training_setup(self, train_dataset: ExtractorDataset, eval_dataset: Optional[ExtractorDataset]):
"""Validate training setup and raise informative errors for edge cases."""
# Check if dataset is empty
if len(train_dataset) == 0:
raise ValueError("Training dataset is empty. Please provide at least one training example.")
# Check if dataset is smaller than batch size
if len(train_dataset) < self.config.batch_size:
logger.warning(
f"Training dataset size ({len(train_dataset)}) is smaller than batch_size "
f"({self.config.batch_size}). Adjusting batch_size to {len(train_dataset)}."
)
# We'll handle this in _create_dataloader by adjusting drop_last
# Check early stopping configuration
if self.config.early_stopping:
if eval_dataset is None:
raise ValueError(
"early_stopping is enabled but no eval_data provided. "
"Please provide eval_data or disable early_stopping."
)
if len(eval_dataset) == 0:
raise ValueError("Evaluation dataset is empty but early_stopping is enabled.")
# Check eval strategy configuration
if self.config.eval_strategy == "steps" and eval_dataset is None:
logger.warning(
"eval_strategy='steps' but no eval_data provided. "
"Evaluation will be skipped."
)
# Warn about very small datasets
if len(train_dataset) < self.config.gradient_accumulation_steps:
logger.warning(
f"Training dataset size ({len(train_dataset)}) is smaller than "
f"gradient_accumulation_steps ({self.config.gradient_accumulation_steps}). "
f"Training may not work as expected."
)
@staticmethod
def _gold_injection_probability(
progress: float, start: float, end: float, hold_fraction: float
) -> float:
"""Hold the initial injection rate, then linearly anneal to the end."""
if progress <= hold_fraction:
return float(start)
fraction = min(
max((progress - hold_fraction) / max(1.0 - hold_fraction, 1e-12), 0.0),
1.0,
)
return float(start + (end - start) * fraction)
@staticmethod
def _soft_iou_anneal_scale(step: int, anneal_steps: int) -> float:
"""Linearly anneal soft-IoU supervision to exact zero."""
if anneal_steps <= 0:
return 0.0
return max(1.0 - step / anneal_steps, 0.0)
@staticmethod
def _ddp_full_touch(model) -> torch.Tensor:
"""An exactly-zero term reaching EVERY trainable parameter on every step, so plain DDP sees the same used set.
Which parameters receive a gradient varies by batch (classifier, count head, relation and record heads each
only on batches carrying their task). Plain DDP (find_unused_parameters=False) then fails "Expected to have
finished reduction"; find_unused_parameters=True cannot be used because forward returns tensors that are not
part of the loss. Touching every parameter every step keeps the used set constant and every gradient reduced."""
return sum(p.sum() * 0.0 for p in model.parameters() if p.requires_grad)
def _backward_one(
self,
batch,
step: int,
use_amp: bool,
amp_dtype,
*,
is_last_micro: bool = True,
) -> torch.Tensor:
"""Run one micro-batch and always enter autograd/DDP collectives."""
if self._skip_counter is None:
self._skip_counter = torch.zeros(
(), dtype=torch.long, device=self.device
)
self._loss_accum = torch.zeros(
(), dtype=torch.float32, device=self.device
)
model = self.model.module if hasattr(self.model, "module") else self.model
boundary_head = getattr(model, "boundary_head", None)
if boundary_head is not None:
planned = max(getattr(self, "_planned_max_steps", 1), 1)
progress = self.global_step / planned
injection_probability = self._gold_injection_probability(
progress,
self.config.gold_injection_start,
self.config.gold_injection_end,
self.config.gold_injection_hold_frac,
)
boundary_head.set_gold_injection_prob(injection_probability)
warmup = getattr(
getattr(boundary_head, "settings", None),
"consistency_warmup_steps",
0,
)
consistency_scale = (
1.0 if warmup <= 0 else min(self.global_step / warmup, 1.0)
)
boundary_head.set_consistency_scale(consistency_scale)
soft_iou_steps = getattr(
getattr(boundary_head, "settings", None),
"soft_iou_anneal_steps",
0,
)
soft_iou_scale = self._soft_iou_anneal_scale(
self.global_step, soft_iou_steps
)
boundary_head.set_soft_iou_scale(soft_iou_scale)
diagnostics_interval = (
self.config.diagnostics_every_n_steps or self.config.logging_steps
)
collect = bool(
boundary_head is not None
and self.config.log_proposal_metrics
and (self.global_step + 1) % diagnostics_interval == 0
)
previous_collect = getattr(boundary_head, "collect_diagnostics", False)
if boundary_head is not None:
boundary_head.collect_diagnostics = collect
sync_ctx = (
contextlib.nullcontext()
if is_last_micro or not self.is_distributed
else self.model.no_sync()
)
try:
with sync_ctx:
with torch.amp.autocast(
device_type=self.device.type,
enabled=use_amp,
dtype=amp_dtype,
):
outputs = self.model(batch)
loss = outputs.get("total_loss")
if loss is None:
zero = getattr(model, "_zero_loss", None)
loss = (
zero(self.device)
if zero is not None
else sum(
(parameter.sum() * 0.0)
for parameter in model.parameters()
if parameter.requires_grad
)
)
elif not loss.requires_grad:
zero = getattr(model, "_zero_loss", None)
touch = (
zero(self.device)
if zero is not None
else sum(
(parameter.sum() * 0.0)
for parameter in model.parameters()
if parameter.requires_grad
)
)
loss = loss.detach() * 0.0 + touch
finite = torch.isfinite(loss.detach()).all()
if self.is_distributed and self.config.ddp_consensus_check:
bad = (~finite).to(dtype=torch.float32)
dist.all_reduce(bad, op=dist.ReduceOp.MAX)
finite = bad == 0
self._loss_finite_flag = finite
loss = torch.where(finite, loss, torch.zeros_like(loss))
reported_loss = loss.detach()
if self.config.gradient_accumulation_steps > 1:
loss = loss / self.config.gradient_accumulation_steps
if self.is_distributed:
loss = loss + self._ddp_full_touch(model)
if self.config.fp16:
self.scaler.scale(loss).backward()
else:
loss.backward()
finally:
if boundary_head is not None:
boundary_head.collect_diagnostics = previous_collect
self._skip_counter.add_((~finite).to(self._skip_counter.dtype))
self._loss_accum.add_(reported_loss.float())
self._last_train_outputs = outputs
return reported_loss
def _renormalize_partial_accumulation(self, micro_batches: int) -> None:
"""Correct gradients from an incomplete accumulation window."""