code/src/train_ssl_dinov3.py · savastakan/hitit-cuneiform-ocr at main

hitit-cuneiform-ocr / code /src /train_ssl_dinov3.py

Initial upload: code + 5 record checkpoints + fuse

f211247 verified 6 months ago

15.4 kB

#!/usr/bin/env python3
"""DINOv3 continual SSL pretraining.
Simplified iBOT/DINO-style training using timm ViT + EMA teacher.
Config-driven (ssl_dinov3_continual.yaml).
Bit-exact BF16, torch.compile, FlashAttention-2.
Data: in_curated_pretrain=True (70K Hitit-ish images)
Output: backbone checkpoint for downstream fine-tune.
"""
import os, sys, json, math, argparse, random, time
from pathlib import Path
from copy import deepcopy

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from PIL import Image
import yaml

ROOT = Path("/arf/scratch/stakan/hitit-proje")

def load_config(path):
    return yaml.safe_load(open(path))

def setup_ddp():
    if 'RANK' in os.environ:
        # Retry CUDA init (transient NVML race on some nodes)
        for attempt in range(3):
            try:
                n = torch.cuda.device_count()
                if n > 0 and torch.cuda.is_available(): break
            except Exception: pass
            time.sleep(2)
        if not torch.cuda.is_available() or torch.cuda.device_count() == 0:
            # Fallback: unset CUDA_VISIBLE_DEVICES and retry
            vis = os.environ.pop('CUDA_VISIBLE_DEVICES', None)
            print(f"[rank {os.environ.get('RANK')}] CUDA init failed; retrying after unset (was {vis})", flush=True)
            import importlib
            importlib.reload(torch.cuda)
            time.sleep(2)
        torch.distributed.init_process_group(backend='nccl')
        local_rank = int(os.environ.get('LOCAL_RANK', 0))
        torch.cuda.set_device(local_rank)
        return True, local_rank, int(os.environ['WORLD_SIZE']), int(os.environ['RANK'])
    return False, 0, 1, 0

def is_main(rank): return rank == 0

def log(msg, rank=0):
    if is_main(rank):
        print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)

class SSLImageDataset(Dataset):
    """Image-only dataset for SSL (no labels).
    Sources: either (a) `in_curated_pretrain=True` flag across all
    datasets/sources/*/manifest.jsonl, or (b) explicit list of manifests
    via manifests_list argument.
    """
    def __init__(self, manifest_path, filter_curated=True, global_size=224, local_size=96,
                 n_local_crops=6, manifests_list=None):
        self.records = []
        import os as _os
        if manifests_list:
            # Explicit list: take every valid image, no in_curated_pretrain filter
            for mf in manifests_list:
                mp = Path(mf)
                if not mp.exists(): continue
                with open(mp) as f:
                    for line in f:
                        r = json.loads(line)
                        if r.get('storage') != 'fs' or not r.get('path'): continue
                        if r.get('integrity_ok') is False: continue
                        if not _os.path.exists(r['path']): continue
                        self.records.append(r['path'])
        else:
            sources_dir = ROOT / 'datasets' / 'sources'
            for src_dir in sources_dir.iterdir():
                if not src_dir.is_dir(): continue
                mf = src_dir / 'manifest.jsonl'
                if not mf.exists(): continue
                with open(mf) as f:
                    for line in f:
                        r = json.loads(line)
                        if filter_curated and not r.get('in_curated_pretrain'):
                            continue
                        if r.get('storage') != 'fs' or not r.get('path'):
                            continue
                        if r.get('integrity_ok') is False:
                            continue
                        if not _os.path.exists(r['path']):
                            continue
                        self.records.append(r['path'])

        self.global_size = global_size
        self.local_size = local_size
        self.n_local_crops = n_local_crops

        # Basic augmentation via torchvision
        from torchvision import transforms
        self.global_transform = transforms.Compose([
            transforms.RandomResizedCrop(global_size, scale=(0.4, 1.0), antialias=True),
            transforms.RandomHorizontalFlip(p=0.0),  # cuneiform yön-duyarlı
            transforms.ColorJitter(0.4, 0.4, 0.2, 0.1),
            transforms.RandomGrayscale(p=0.2),
            transforms.ToTensor(),
            transforms.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
        ])
        self.local_transform = transforms.Compose([
            transforms.RandomResizedCrop(local_size, scale=(0.05, 0.4), antialias=True),
            transforms.ColorJitter(0.4, 0.4, 0.2, 0.1),
            transforms.ToTensor(),
            transforms.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
        ])

    def __len__(self): return len(self.records)

    def __getitem__(self, idx):
        try:
            img = Image.open(self.records[idx]).convert('RGB')
        except Exception:
            img = Image.new('RGB', (self.global_size, self.global_size))
        globals_ = [self.global_transform(img) for _ in range(2)]
        locals_ = [self.local_transform(img) for _ in range(self.n_local_crops)]
        return globals_, locals_

def collate_fn(batch):
    # batch: list of ([g1, g2], [l1..l6])
    globals_ = torch.stack([g for item in batch for g in item[0]])  # 2B × 3 × H × W
    locals_ = torch.stack([l for item in batch for l in item[1]])   # nB × 3 × h × w
    return globals_, locals_

class ProjectionHead(nn.Module):
    """DINO projection head.
    weight_norm + last_layer.weight_g.fill_(1) ile birlikte F.normalize uygulanması
    çıkışı sabit-norm vektörlere kilitliyor ve teacher==student kalıyordu (loss
    uniform entropy ln(65536)≈11.09'da donuyor). last_layer'ı normal Linear yapıp
    init'i küçük tutuyoruz; bottleneck normalize-edilmiş kalsın.
    """
    def __init__(self, in_dim=768, out_dim=65536, hidden_dim=2048, bottleneck=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, hidden_dim), nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim), nn.GELU(),
            nn.Linear(hidden_dim, bottleneck),
        )
        self.last_layer = nn.Linear(bottleneck, out_dim, bias=False)
        nn.init.trunc_normal_(self.last_layer.weight, std=0.02)
    def forward(self, x):
        x = self.mlp(x)
        x = F.normalize(x, dim=-1)
        return self.last_layer(x)

class DINOLoss(nn.Module):
    def __init__(self, out_dim=65536, student_temp=0.1, teacher_temp=0.04,
                 center_momentum=0.9):
        super().__init__()
        self.student_temp = student_temp
        self.teacher_temp = teacher_temp
        self.center_momentum = center_momentum
        self.register_buffer('center', torch.zeros(1, out_dim))

    def forward(self, student_out, teacher_out):
        """
        student_out: (N_total_crops × B, out_dim) — 2 global + K local
        teacher_out: (2 × B, out_dim) — only global
        """
        s = student_out / self.student_temp
        t = F.softmax((teacher_out - self.center) / self.teacher_temp, dim=-1).detach()

        # Count local + global
        B = t.size(0) // 2
        total_crops = s.size(0) // B
        s = s.chunk(total_crops)
        t = t.chunk(2)

        total_loss = 0
        n_loss_terms = 0
        for iq, q in enumerate(t):
            for v in range(len(s)):
                if v == iq: continue  # skip same view
                loss = torch.sum(-q * F.log_softmax(s[v], dim=-1), dim=-1)
                total_loss += loss.mean()
                n_loss_terms += 1
        total_loss = total_loss / max(n_loss_terms, 1)
        self.update_center(teacher_out)
        return total_loss

    @torch.no_grad()
    def update_center(self, teacher_out):
        batch_center = teacher_out.mean(0, keepdim=True)
        if torch.distributed.is_initialized():
            torch.distributed.all_reduce(batch_center)
            batch_center /= torch.distributed.get_world_size()
        self.center = self.center * self.center_momentum + batch_center * (1 - self.center_momentum)

def build_backbone(model_name='vit_base_patch14_dinov2', pretrained=True):
    import timm
    # dynamic_img_size=True → ViT accepts any resolution (224 global + 96 local)
    try:
        model = timm.create_model(model_name, pretrained=pretrained, num_classes=0,
                                   img_size=224, dynamic_img_size=True)
    except Exception:
        try:
            model = timm.create_model('vit_base_patch16_224', pretrained=pretrained,
                                       num_classes=0, dynamic_img_size=True)
        except Exception:
            # Son çare: sabit image_size, locals disable
            model = timm.create_model('vit_base_patch16_224', pretrained=pretrained, num_classes=0)
    return model

@torch.no_grad()
def ema_update(teacher, student, momentum):
    for tp, sp in zip(teacher.parameters(), student.parameters()):
        tp.data.mul_(momentum).add_(sp.data, alpha=1 - momentum)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--config', default=str(ROOT / 'hitit_ocr/configs/ssl_dinov3_continual.yaml'))
    ap.add_argument('--output', default=str(ROOT / 'hitit_ocr/runs/ssl_dinov3_continual/'))
    ap.add_argument('--ddp-world-size', type=int, default=0)
    args = ap.parse_args()

    cfg = load_config(args.config)
    is_ddp, local_rank, world_size, global_rank = setup_ddp()
    device = torch.device(f'cuda:{local_rank}' if torch.cuda.is_available() else 'cpu')

    log(f"DDP: {is_ddp}, world_size={world_size}, rank={global_rank}", global_rank)
    log(f"Device: {device}", global_rank)

    Path(args.output).mkdir(parents=True, exist_ok=True)

    # Data
    crops = cfg.get('crops', {})
    manifest = ROOT / 'datasets/unified/classification/manifest.jsonl'
    data_cfg = cfg.get('data', {})
    manifests_list = data_cfg.get('manifests')  # if set in config, use them
    filter_curated = data_cfg.get('filter') is not None  # None or 'null' → all
    dataset = SSLImageDataset(
        manifest, filter_curated=filter_curated,
        global_size=crops.get('global_size', 224),
        local_size=crops.get('local_size', 96),
        n_local_crops=crops.get('local_count', 6),
        manifests_list=manifests_list,
    )
    log(f"Dataset: {len(dataset)} images", global_rank)

    sampler = DistributedSampler(dataset) if is_ddp else None
    batch_size = cfg.get('batch_size', 256) // max(world_size, 1)
    loader = DataLoader(
        dataset, batch_size=batch_size, shuffle=(sampler is None),
        sampler=sampler, num_workers=8, pin_memory=True, drop_last=True,
        collate_fn=collate_fn, persistent_workers=True,
    )

    # Models
    backbone_name = cfg.get('model', {}).get('teacher', 'vit_base_patch14_dinov2')
    student_bb = build_backbone(backbone_name).to(device)
    teacher_bb = build_backbone(backbone_name).to(device)
    teacher_bb.load_state_dict(student_bb.state_dict())
    for p in teacher_bb.parameters(): p.requires_grad = False

    embed_dim = getattr(student_bb, 'num_features', 768)
    student_head = ProjectionHead(in_dim=embed_dim, out_dim=65536).to(device)
    teacher_head = ProjectionHead(in_dim=embed_dim, out_dim=65536).to(device)
    teacher_head.load_state_dict(student_head.state_dict())
    for p in teacher_head.parameters(): p.requires_grad = False

    if is_ddp:
        student_bb = DDP(student_bb, device_ids=[local_rank])
        student_head = DDP(student_head, device_ids=[local_rank])

    # Loss
    loss_fn = DINOLoss().to(device)

    # Optimizer
    params = [
        {'params': (student_bb.module if is_ddp else student_bb).parameters()},
        {'params': (student_head.module if is_ddp else student_head).parameters()},
    ]
    lr = cfg.get('optimizer', {}).get('lr', 1e-4)
    wd = cfg.get('optimizer', {}).get('weight_decay', 0.04)
    optimizer = torch.optim.AdamW(params, lr=lr, weight_decay=wd, fused=True)

    # BF16 autocast
    use_bf16 = cfg.get('use_bf16', True)
    dtype = torch.bfloat16 if use_bf16 else torch.float32

    epochs = cfg.get('epochs', 20)
    ema_start = cfg.get('ema', {}).get('teacher_momentum_start', 0.996)

    log(f"Training: {epochs} epochs, batch={batch_size}/GPU × {world_size} GPUs", global_rank)
    log(f"BF16: {use_bf16}", global_rank)

    step = 0
    total_steps = len(loader) * epochs
    t0 = time.time()

    for epoch in range(epochs):
        if sampler: sampler.set_epoch(epoch)
        epoch_loss = 0.0
        n_batches = 0

        for globals_, locals_ in loader:
            globals_ = globals_.to(device, non_blocking=True)
            locals_ = locals_.to(device, non_blocking=True)

            with torch.amp.autocast('cuda', dtype=dtype, enabled=use_bf16):
                # Teacher: only globals (224×224)
                with torch.no_grad():
                    t_feat = teacher_bb(globals_)
                    t_out = teacher_head(t_feat)

                # Student: globals AND locals separately (different sizes)
                s_feat_g = student_bb(globals_)
                s_feat_l = student_bb(locals_)
                s_out_g = student_head(s_feat_g)
                s_out_l = student_head(s_feat_l)
                s_out = torch.cat([s_out_g, s_out_l], dim=0)

                loss = loss_fn(s_out, t_out)

            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(
                (student_bb.module if is_ddp else student_bb).parameters(), 3.0)
            optimizer.step()

            # EMA update
            m = ema_start + (1.0 - ema_start) * step / total_steps
            ema_update(teacher_bb, student_bb.module if is_ddp else student_bb, m)
            ema_update(teacher_head, student_head.module if is_ddp else student_head, m)

            epoch_loss += loss.item()
            n_batches += 1
            step += 1

            if step % 50 == 0:
                rate = step / (time.time() - t0)
                log(f"  epoch={epoch} step={step}/{total_steps} loss={loss.item():.4f} rate={rate:.2f}/s", global_rank)

        log(f"Epoch {epoch}: avg_loss={epoch_loss/max(1,n_batches):.4f}", global_rank)

        if is_main(global_rank) and (epoch + 1) % cfg.get('save_every_epochs', 5) == 0:
            ckpt = {
                'epoch': epoch,
                'backbone': (teacher_bb.state_dict()),  # use teacher for downstream
                'head': teacher_head.state_dict(),
                'cfg': cfg,
            }
            torch.save(ckpt, Path(args.output) / f'checkpoint_e{epoch:02d}.pt')
            torch.save(ckpt, Path(args.output) / 'checkpoint.pt')
            log(f"Saved checkpoint epoch {epoch}", global_rank)

    # Final
    if is_main(global_rank):
        ckpt = {
            'epoch': epochs,
            'backbone': (teacher_bb.state_dict()),
            'head': teacher_head.state_dict(),
            'cfg': cfg,
        }
        torch.save(ckpt, Path(args.output) / 'checkpoint.pt')
        log(f"DONE: final checkpoint saved → {args.output}/checkpoint.pt", global_rank)

    if is_ddp:
        torch.distributed.destroy_process_group()

if __name__ == '__main__':
    main()

	#!/usr/bin/env python3
	"""DINOv3 continual SSL pretraining.

	Simplified iBOT/DINO-style training using timm ViT + EMA teacher.
	Config-driven (ssl_dinov3_continual.yaml).
	Bit-exact BF16, torch.compile, FlashAttention-2.

	Data: in_curated_pretrain=True (70K Hitit-ish images)
	Output: backbone checkpoint for downstream fine-tune.
	"""
	import os, sys, json, math, argparse, random, time
	from pathlib import Path
	from copy import deepcopy

	import numpy as np
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from torch.utils.data import Dataset, DataLoader, DistributedSampler
	from torch.nn.parallel import DistributedDataParallel as DDP
	from PIL import Image
	import yaml

	ROOT = Path("/arf/scratch/stakan/hitit-proje")

	def load_config(path):
	return yaml.safe_load(open(path))

	def setup_ddp():
	if 'RANK' in os.environ:
	# Retry CUDA init (transient NVML race on some nodes)
	for attempt in range(3):
	try:
	n = torch.cuda.device_count()
	if n > 0 and torch.cuda.is_available(): break
	except Exception: pass
	time.sleep(2)
	if not torch.cuda.is_available() or torch.cuda.device_count() == 0:
	# Fallback: unset CUDA_VISIBLE_DEVICES and retry
	vis = os.environ.pop('CUDA_VISIBLE_DEVICES', None)
	print(f"[rank {os.environ.get('RANK')}] CUDA init failed; retrying after unset (was {vis})", flush=True)
	import importlib
	importlib.reload(torch.cuda)
	time.sleep(2)
	torch.distributed.init_process_group(backend='nccl')
	local_rank = int(os.environ.get('LOCAL_RANK', 0))
	torch.cuda.set_device(local_rank)
	return True, local_rank, int(os.environ['WORLD_SIZE']), int(os.environ['RANK'])
	return False, 0, 1, 0

	def is_main(rank): return rank == 0

	def log(msg, rank=0):
	if is_main(rank):
	print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)

	class SSLImageDataset(Dataset):
	"""Image-only dataset for SSL (no labels).

	Sources: either (a) `in_curated_pretrain=True` flag across all
	datasets/sources/*/manifest.jsonl, or (b) explicit list of manifests
	via manifests_list argument.
	"""
	def __init__(self, manifest_path, filter_curated=True, global_size=224, local_size=96,
	n_local_crops=6, manifests_list=None):
	self.records = []
	import os as _os
	if manifests_list:
	# Explicit list: take every valid image, no in_curated_pretrain filter
	for mf in manifests_list:
	mp = Path(mf)
	if not mp.exists(): continue
	with open(mp) as f:
	for line in f:
	r = json.loads(line)
	if r.get('storage') != 'fs' or not r.get('path'): continue
	if r.get('integrity_ok') is False: continue
	if not _os.path.exists(r['path']): continue
	self.records.append(r['path'])
	else:
	sources_dir = ROOT / 'datasets' / 'sources'
	for src_dir in sources_dir.iterdir():
	if not src_dir.is_dir(): continue
	mf = src_dir / 'manifest.jsonl'
	if not mf.exists(): continue
	with open(mf) as f:
	for line in f:
	r = json.loads(line)
	if filter_curated and not r.get('in_curated_pretrain'):
	continue
	if r.get('storage') != 'fs' or not r.get('path'):
	continue
	if r.get('integrity_ok') is False:
	continue
	if not _os.path.exists(r['path']):
	continue
	self.records.append(r['path'])

	self.global_size = global_size
	self.local_size = local_size
	self.n_local_crops = n_local_crops

	# Basic augmentation via torchvision
	from torchvision import transforms
	self.global_transform = transforms.Compose([
	transforms.RandomResizedCrop(global_size, scale=(0.4, 1.0), antialias=True),
	transforms.RandomHorizontalFlip(p=0.0), # cuneiform yön-duyarlı
	transforms.ColorJitter(0.4, 0.4, 0.2, 0.1),
	transforms.RandomGrayscale(p=0.2),
	transforms.ToTensor(),
	transforms.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
	])
	self.local_transform = transforms.Compose([
	transforms.RandomResizedCrop(local_size, scale=(0.05, 0.4), antialias=True),
	transforms.ColorJitter(0.4, 0.4, 0.2, 0.1),
	transforms.ToTensor(),
	transforms.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
	])

	def __len__(self): return len(self.records)

	def __getitem__(self, idx):
	try:
	img = Image.open(self.records[idx]).convert('RGB')
	except Exception:
	img = Image.new('RGB', (self.global_size, self.global_size))
	globals_ = [self.global_transform(img) for _ in range(2)]
	locals_ = [self.local_transform(img) for _ in range(self.n_local_crops)]
	return globals_, locals_

	def collate_fn(batch):
	# batch: list of ([g1, g2], [l1..l6])
	globals_ = torch.stack([g for item in batch for g in item[0]]) # 2B × 3 × H × W
	locals_ = torch.stack([l for item in batch for l in item[1]]) # nB × 3 × h × w
	return globals_, locals_

	class ProjectionHead(nn.Module):
	"""DINO projection head.

	weight_norm + last_layer.weight_g.fill_(1) ile birlikte F.normalize uygulanması
	çıkışı sabit-norm vektörlere kilitliyor ve teacher==student kalıyordu (loss
	uniform entropy ln(65536)≈11.09'da donuyor). last_layer'ı normal Linear yapıp
	init'i küçük tutuyoruz; bottleneck normalize-edilmiş kalsın.
	"""
	def __init__(self, in_dim=768, out_dim=65536, hidden_dim=2048, bottleneck=256):
	super().__init__()
	self.mlp = nn.Sequential(
	nn.Linear(in_dim, hidden_dim), nn.GELU(),
	nn.Linear(hidden_dim, hidden_dim), nn.GELU(),
	nn.Linear(hidden_dim, bottleneck),
	)
	self.last_layer = nn.Linear(bottleneck, out_dim, bias=False)
	nn.init.trunc_normal_(self.last_layer.weight, std=0.02)
	def forward(self, x):
	x = self.mlp(x)
	x = F.normalize(x, dim=-1)
	return self.last_layer(x)

	class DINOLoss(nn.Module):
	def __init__(self, out_dim=65536, student_temp=0.1, teacher_temp=0.04,
	center_momentum=0.9):
	super().__init__()
	self.student_temp = student_temp
	self.teacher_temp = teacher_temp
	self.center_momentum = center_momentum
	self.register_buffer('center', torch.zeros(1, out_dim))

	def forward(self, student_out, teacher_out):
	"""
	student_out: (N_total_crops × B, out_dim) — 2 global + K local
	teacher_out: (2 × B, out_dim) — only global
	"""
	s = student_out / self.student_temp
	t = F.softmax((teacher_out - self.center) / self.teacher_temp, dim=-1).detach()

	# Count local + global
	B = t.size(0) // 2
	total_crops = s.size(0) // B
	s = s.chunk(total_crops)
	t = t.chunk(2)

	total_loss = 0
	n_loss_terms = 0
	for iq, q in enumerate(t):
	for v in range(len(s)):
	if v == iq: continue # skip same view
	loss = torch.sum(-q * F.log_softmax(s[v], dim=-1), dim=-1)
	total_loss += loss.mean()
	n_loss_terms += 1
	total_loss = total_loss / max(n_loss_terms, 1)
	self.update_center(teacher_out)
	return total_loss

	@torch.no_grad()
	def update_center(self, teacher_out):
	batch_center = teacher_out.mean(0, keepdim=True)
	if torch.distributed.is_initialized():
	torch.distributed.all_reduce(batch_center)
	batch_center /= torch.distributed.get_world_size()
	self.center = self.center * self.center_momentum + batch_center * (1 - self.center_momentum)

	def build_backbone(model_name='vit_base_patch14_dinov2', pretrained=True):
	import timm
	# dynamic_img_size=True → ViT accepts any resolution (224 global + 96 local)
	try:
	model = timm.create_model(model_name, pretrained=pretrained, num_classes=0,
	img_size=224, dynamic_img_size=True)
	except Exception:
	try:
	model = timm.create_model('vit_base_patch16_224', pretrained=pretrained,
	num_classes=0, dynamic_img_size=True)
	except Exception:
	# Son çare: sabit image_size, locals disable
	model = timm.create_model('vit_base_patch16_224', pretrained=pretrained, num_classes=0)
	return model

	@torch.no_grad()
	def ema_update(teacher, student, momentum):
	for tp, sp in zip(teacher.parameters(), student.parameters()):
	tp.data.mul_(momentum).add_(sp.data, alpha=1 - momentum)

	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument('--config', default=str(ROOT / 'hitit_ocr/configs/ssl_dinov3_continual.yaml'))
	ap.add_argument('--output', default=str(ROOT / 'hitit_ocr/runs/ssl_dinov3_continual/'))
	ap.add_argument('--ddp-world-size', type=int, default=0)
	args = ap.parse_args()

	cfg = load_config(args.config)
	is_ddp, local_rank, world_size, global_rank = setup_ddp()
	device = torch.device(f'cuda:{local_rank}' if torch.cuda.is_available() else 'cpu')

	log(f"DDP: {is_ddp}, world_size={world_size}, rank={global_rank}", global_rank)
	log(f"Device: {device}", global_rank)

	Path(args.output).mkdir(parents=True, exist_ok=True)

	# Data
	crops = cfg.get('crops', {})
	manifest = ROOT / 'datasets/unified/classification/manifest.jsonl'
	data_cfg = cfg.get('data', {})
	manifests_list = data_cfg.get('manifests') # if set in config, use them
	filter_curated = data_cfg.get('filter') is not None # None or 'null' → all
	dataset = SSLImageDataset(
	manifest, filter_curated=filter_curated,
	global_size=crops.get('global_size', 224),
	local_size=crops.get('local_size', 96),
	n_local_crops=crops.get('local_count', 6),
	manifests_list=manifests_list,
	)
	log(f"Dataset: {len(dataset)} images", global_rank)

	sampler = DistributedSampler(dataset) if is_ddp else None
	batch_size = cfg.get('batch_size', 256) // max(world_size, 1)
	loader = DataLoader(
	dataset, batch_size=batch_size, shuffle=(sampler is None),
	sampler=sampler, num_workers=8, pin_memory=True, drop_last=True,
	collate_fn=collate_fn, persistent_workers=True,
	)

	# Models
	backbone_name = cfg.get('model', {}).get('teacher', 'vit_base_patch14_dinov2')
	student_bb = build_backbone(backbone_name).to(device)
	teacher_bb = build_backbone(backbone_name).to(device)
	teacher_bb.load_state_dict(student_bb.state_dict())
	for p in teacher_bb.parameters(): p.requires_grad = False

	embed_dim = getattr(student_bb, 'num_features', 768)
	student_head = ProjectionHead(in_dim=embed_dim, out_dim=65536).to(device)
	teacher_head = ProjectionHead(in_dim=embed_dim, out_dim=65536).to(device)
	teacher_head.load_state_dict(student_head.state_dict())
	for p in teacher_head.parameters(): p.requires_grad = False

	if is_ddp:
	student_bb = DDP(student_bb, device_ids=[local_rank])
	student_head = DDP(student_head, device_ids=[local_rank])

	# Loss
	loss_fn = DINOLoss().to(device)

	# Optimizer
	params = [
	{'params': (student_bb.module if is_ddp else student_bb).parameters()},
	{'params': (student_head.module if is_ddp else student_head).parameters()},
	]
	lr = cfg.get('optimizer', {}).get('lr', 1e-4)
	wd = cfg.get('optimizer', {}).get('weight_decay', 0.04)
	optimizer = torch.optim.AdamW(params, lr=lr, weight_decay=wd, fused=True)

	# BF16 autocast
	use_bf16 = cfg.get('use_bf16', True)
	dtype = torch.bfloat16 if use_bf16 else torch.float32

	epochs = cfg.get('epochs', 20)
	ema_start = cfg.get('ema', {}).get('teacher_momentum_start', 0.996)

	log(f"Training: {epochs} epochs, batch={batch_size}/GPU × {world_size} GPUs", global_rank)
	log(f"BF16: {use_bf16}", global_rank)

	step = 0
	total_steps = len(loader) * epochs
	t0 = time.time()

	for epoch in range(epochs):
	if sampler: sampler.set_epoch(epoch)
	epoch_loss = 0.0
	n_batches = 0

	for globals_, locals_ in loader:
	globals_ = globals_.to(device, non_blocking=True)
	locals_ = locals_.to(device, non_blocking=True)

	with torch.amp.autocast('cuda', dtype=dtype, enabled=use_bf16):
	# Teacher: only globals (224×224)
	with torch.no_grad():
	t_feat = teacher_bb(globals_)
	t_out = teacher_head(t_feat)

	# Student: globals AND locals separately (different sizes)
	s_feat_g = student_bb(globals_)
	s_feat_l = student_bb(locals_)
	s_out_g = student_head(s_feat_g)
	s_out_l = student_head(s_feat_l)
	s_out = torch.cat([s_out_g, s_out_l], dim=0)

	loss = loss_fn(s_out, t_out)

	optimizer.zero_grad(set_to_none=True)
	loss.backward()
	torch.nn.utils.clip_grad_norm_(
	(student_bb.module if is_ddp else student_bb).parameters(), 3.0)
	optimizer.step()

	# EMA update
	m = ema_start + (1.0 - ema_start) * step / total_steps
	ema_update(teacher_bb, student_bb.module if is_ddp else student_bb, m)
	ema_update(teacher_head, student_head.module if is_ddp else student_head, m)

	epoch_loss += loss.item()
	n_batches += 1
	step += 1

	if step % 50 == 0:
	rate = step / (time.time() - t0)
	log(f" epoch={epoch} step={step}/{total_steps} loss={loss.item():.4f} rate={rate:.2f}/s", global_rank)

	log(f"Epoch {epoch}: avg_loss={epoch_loss/max(1,n_batches):.4f}", global_rank)

	if is_main(global_rank) and (epoch + 1) % cfg.get('save_every_epochs', 5) == 0:
	ckpt = {
	'epoch': epoch,
	'backbone': (teacher_bb.state_dict()), # use teacher for downstream
	'head': teacher_head.state_dict(),
	'cfg': cfg,
	}
	torch.save(ckpt, Path(args.output) / f'checkpoint_e{epoch:02d}.pt')
	torch.save(ckpt, Path(args.output) / 'checkpoint.pt')
	log(f"Saved checkpoint epoch {epoch}", global_rank)

	# Final
	if is_main(global_rank):
	ckpt = {
	'epoch': epochs,
	'backbone': (teacher_bb.state_dict()),
	'head': teacher_head.state_dict(),
	'cfg': cfg,
	}
	torch.save(ckpt, Path(args.output) / 'checkpoint.pt')
	log(f"DONE: final checkpoint saved → {args.output}/checkpoint.pt", global_rank)

	if is_ddp:
	torch.distributed.destroy_process_group()

	if __name__ == '__main__':
	main()