pixr 0.1.0.dev0__py3-none-any.whl

pixr/__init__.py

@@ -0,0 +1,3 @@
+from .models import Cruller, ModelCfg, ImageEncoderCfg, TextDecoderCfg, get_model_config, list_models
+from .tokenizers import TokenizerCfg
+from .version import __version__

pixr/app/eval.py

@@ -0,0 +1,95 @@
+""" pixr evaluation app
+
+Example (DocVQA validation ANLS):
+    python -m pixr.app.eval \\
+        --task docvqa \\
+        --eval-data.source pixparse/docvqa-single-page-questions --eval-data.split validation \\
+        --eval-data.format hfds --eval-data.batch-size 32 \\
+        --checkpoint-path output/<experiment>/checkpoints/checkpoint-latest.pt
+"""
+import json
+import logging
+import os
+import re
+import sys
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+from chugr.common import DataArg, Serializable
+
+from pixr.data import create_loader
+from pixr.framework import DeviceEnv, destroy_distributed, setup_logging, random_seed
+from pixr.task import get_task, get_task_cfgs
+from pixr.utils import load_checkpoint
+from pixr.utils.cli import parse_cfg
+
+_logger = logging.getLogger('eval')
+
+
+@dataclass
+class EvalCfg(Serializable):
+    eval_data: DataArg = field(default_factory=lambda: DataArg(source=''))
+    task_name: str = 'docvqa'  # selected via --task <name> (or task_name in a --config file)
+    task: Any = None  # task config, type set by the selected task
+    checkpoint_path: str = ''
+    use_ema: bool = False  # evaluate EMA weights if present in the checkpoint
+    output_dir: Optional[str] = None  # default: <checkpoint dir>/eval
+    seed: int = 42
+
+
+def _run(argv=None):
+    eval_cfg: EvalCfg = parse_cfg(
+        EvalCfg,
+        argv=argv,
+        task_cfgs=get_task_cfgs(),
+        default_task='docvqa',
+        description='pixr eval',
+    )
+    task_name = eval_cfg.task_name
+
+    device_env = DeviceEnv()
+    random_seed(eval_cfg.seed, rank=device_env.global_rank)
+    setup_logging(None)
+    if device_env.is_primary():
+        _logger.info('Eval config:\n' + eval_cfg.dumps_yaml())
+
+    assert eval_cfg.checkpoint_path, 'A checkpoint is required for evaluation (--checkpoint-path).'
+    task = get_task(task_name)(eval_cfg.task, device_env=device_env, checkpoint_path=eval_cfg.checkpoint_path)
+    if eval_cfg.use_ema:
+        checkpoint = load_checkpoint(eval_cfg.checkpoint_path)
+        assert 'model_ema' in checkpoint, 'Checkpoint has no EMA weights.'
+        task.trainable_module.load_state_dict(checkpoint['model_ema'])
+    task.to(device_env.device).eval()
+
+    loader = create_loader(
+        eval_cfg.eval_data,
+        task.eval_data_cfg(),
+        is_training=False,
+        world_size=device_env.world_size,
+        global_rank=device_env.global_rank,
+        local_rank=device_env.local_rank,
+    )
+    metrics = task.evaluate(loader.loader)
+
+    if device_env.is_primary():
+        _logger.info(f'{task_name} results: {metrics}')
+        output_dir = eval_cfg.output_dir or os.path.join(
+            os.path.dirname(os.path.abspath(eval_cfg.checkpoint_path)), 'eval')
+        split = re.sub(r'[^A-Za-z0-9_.+-]+', '_', eval_cfg.eval_data.split or 'eval').strip('_')
+        name = f'{task_name}-{split}'
+        task.save_results(output_dir, name=name)
+        with open(os.path.join(output_dir, f'{name}_metrics.json'), 'w') as f:
+            json.dump(metrics, f, indent=2)
+        _logger.info(f'Results written to {output_dir}')
+
+
+def main(argv=None):
+    """ Entry point for both launch forms (`pixr-eval` console script, `python -m pixr.app.eval`). """
+    code = _run(argv)
+    # NOTE only after a normal return, w/ an exception other ranks may still be in collectives, the launcher tears down
+    destroy_distributed()
+    return code
+
+
+if __name__ == '__main__':
+    sys.exit(main())

pixr/app/train.py

@@ -0,0 +1,192 @@
+""" pixr training app
+
+Example (read pretraining on IDL streamed from the HF hub, 2 GPUs):
+    python -m torch.distributed.run --nproc_per_node 2 -m pixr.app.train \\
+        --task read --task.model.name cruller_naflex_siglip2_base \\
+        --train-data.source "pipe:curl -s -f -L https://huggingface.co/datasets/pixparse/idl-wds/resolve/main/idl-train-0{0000..2999}.tar" \\
+        --train-data.num-samples 1000000 --train-data.batch-size 16 --train-data.num-workers 16 \\
+        --trainer.opt.learning-rate 2e-4
+"""
+import logging
+import os
+import sys
+from dataclasses import dataclass, field, replace
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import torch
+
+from chugr.common import DataArg, Serializable
+
+from pixr.data import create_loader
+from pixr.framework import DeviceEnv, destroy_distributed, Monitor, Trainer, TrainerCfg, setup_logging, random_seed
+from pixr.task import get_task, get_task_cfgs
+from pixr.utils import clean_name, get_latest_checkpoint
+from pixr.utils.cli import parse_cfg
+
+_logger = logging.getLogger('train')
+
+
+@dataclass
+class TrainCfg(Serializable):
+    train_data: DataArg = field(default_factory=lambda: DataArg(source=''))
+    eval_data: Optional[DataArg] = None
+    task_name: str = 'read'  # selected via --task <name> (or task_name in a --config file)
+    task: Any = None  # task config, type set by the selected task
+    trainer: TrainerCfg = field(default_factory=TrainerCfg)
+
+    experiment: Optional[str] = None  # experiment name, auto-generated if None
+    output_dir: str = './output'
+    log_filename: str = 'out.log'
+    resume: str = ''  # resume training from checkpoint w/ full train state ('latest' for most recent)
+    checkpoint_path: str = ''  # initialize model weights (and tokens / arch) from this checkpoint
+    eval_frequency: int = 1  # evaluate every N intervals (if eval_data set)
+    eval_ema: bool = False  # evaluate the EMA weights (if trainer.ema_decay set)
+    checkpoint_frequency: int = 1  # save checkpoint every N intervals (last interval always saved)
+    checkpoint_keep: int = 3  # number of most recent interval checkpoints to keep (0 == keep all)
+    seed: int = 42
+
+    wandb: bool = False
+    wandb_project: str = 'pixr'
+    tensorboard: bool = False
+    log_eval_data: bool = False
+
+
+def _save_checkpoint(state: Dict[str, Any], checkpoint_dir: str, interval: int, keep: int):
+    path = os.path.join(checkpoint_dir, f'checkpoint-{interval}.pt')
+    tmp_path = path + '.tmp'
+    torch.save(state, tmp_path)
+    os.replace(tmp_path, path)
+    latest = os.path.join(checkpoint_dir, 'checkpoint-latest.pt')
+    if os.path.lexists(latest):
+        os.remove(latest)
+    os.symlink(os.path.basename(path), latest)
+    if keep > 0:
+        for f in os.listdir(checkpoint_dir):
+            stem = f[len('checkpoint-'):-3]
+            if f.startswith('checkpoint-') and stem.isdigit() and int(stem) <= interval - keep:
+                os.remove(os.path.join(checkpoint_dir, f))
+    return path
+
+
+def _run(argv=None):
+    train_cfg: TrainCfg = parse_cfg(
+        TrainCfg,
+        argv=argv,
+        task_cfgs=get_task_cfgs(),
+        default_task='read',
+        description='pixr train',
+    )
+    task_name = train_cfg.task_name
+
+    device_env = DeviceEnv()
+    random_seed(train_cfg.seed, rank=device_env.global_rank)
+
+    if train_cfg.experiment is None:
+        date_str = datetime.now().strftime("%Y%m%d-%H%M%S")
+        if device_env.world_size > 1:
+            date_str = device_env.broadcast_object(date_str)
+        experiment = '-'.join([
+            date_str,
+            f"task_{task_name}",
+            f"model_{clean_name(train_cfg.task.model.name)}",
+            f"lr_{'{:.1e}'.format(train_cfg.trainer.opt.learning_rate)}",
+            f"b_{train_cfg.train_data.batch_size}",
+        ])
+        train_cfg = replace(train_cfg, experiment=experiment)
+
+    experiment_path = os.path.join(train_cfg.output_dir, train_cfg.experiment)
+    checkpoint_dir = os.path.join(experiment_path, 'checkpoints')
+    resume = train_cfg.resume
+    if resume == 'latest':
+        resume = get_latest_checkpoint(checkpoint_dir) if os.path.isdir(checkpoint_dir) else ''
+        resume = resume or ''
+    log_path = None
+    exists = False
+    if device_env.is_primary():
+        log_path = os.path.join(experiment_path, train_cfg.log_filename)
+        exists = os.path.exists(log_path) and not resume
+        if exists:
+            print(f"Experiment {train_cfg.experiment} already exists, set --resume or a new --experiment.")
+        else:
+            os.makedirs(checkpoint_dir, exist_ok=True)
+            train_cfg.save(os.path.join(experiment_path, 'config.yaml'))
+    if device_env.world_size > 1:
+        exists = device_env.broadcast_object(exists)  # all ranks exit together
+    if exists:
+        return -1
+
+    setup_logging(log_path if device_env.is_primary() else None)
+    if not device_env.is_primary():
+        logging.root.setLevel(logging.WARNING)
+    _logger.info(f"Device env is {device_env}")
+    if device_env.is_primary():
+        _logger.info('Train config:\n' + train_cfg.dumps_yaml())
+
+    monitor = Monitor(
+        train_cfg.experiment,
+        output_dir=experiment_path,
+        wandb=train_cfg.wandb,
+        wandb_project=train_cfg.wandb_project,
+        tensorboard=train_cfg.tensorboard,
+        output_enabled=device_env.is_primary(),
+        log_eval_data=train_cfg.log_eval_data,
+        hparams=train_cfg.to_dict(),
+    )
+
+    task = get_task(task_name)(train_cfg.task, device_env=device_env, checkpoint_path=train_cfg.checkpoint_path)
+    trainer = Trainer(task, train_cfg.trainer, device_env=device_env, monitor=monitor)
+
+    assert train_cfg.train_data.source, 'Train dataset (--train-data.source) must be set.'
+    dist_kwargs = dict(
+        world_size=device_env.world_size,
+        global_rank=device_env.global_rank,
+        local_rank=device_env.local_rank,
+    )
+    train_loader = create_loader(
+        train_cfg.train_data, task.train_data_cfg(), is_training=True, seed=train_cfg.seed, **dist_kwargs)
+    eval_loader = None
+    if train_cfg.eval_data is not None and train_cfg.eval_data.source:
+        eval_loader = create_loader(train_cfg.eval_data, task.eval_data_cfg(), is_training=False, **dist_kwargs)
+
+    trainer.setup(
+        num_batches_per_interval=train_loader.num_batches,
+        resume=resume,
+        mean_batch_size=train_loader.mean_batch_size,
+    )
+    if device_env.is_primary():
+        num_params = sum(p.numel() for p in task.trainable_module.parameters())
+        _logger.info(f'Model params: {num_params / 1e6:.1f}M, steps per interval: {trainer.num_steps_per_interval}')
+
+    for interval in range(trainer.interval_idx, train_cfg.trainer.num_intervals):
+        train_loader.set_interval(interval)
+        results = {'train': trainer.train_interval(train_loader)}
+
+        last = interval + 1 == train_cfg.trainer.num_intervals
+        if eval_loader is not None and ((interval + 1) % train_cfg.eval_frequency == 0 or last):
+            # w/ FSDP, params are gathered for eval so each rank decodes its eval shard independently
+            with task.full_params(use_ema=train_cfg.eval_ema):
+                eval_metrics = task.evaluate(eval_loader.loader, use_ema=train_cfg.eval_ema)
+            results['eval'] = eval_metrics
+            monitor.log_phase('eval', interval=interval, **eval_metrics)
+
+        if device_env.is_primary():
+            monitor.write_summary(results, index=interval)
+        if (interval + 1) % train_cfg.checkpoint_frequency == 0 or last:
+            state = trainer.state_dict()  # collective w/ FSDP
+            if device_env.is_primary():
+                path = _save_checkpoint(state, checkpoint_dir, interval, train_cfg.checkpoint_keep)
+                _logger.info(f'Saved checkpoint {path}')
+            del state
+
+
+def main(argv=None):
+    """ Entry point for both launch forms (`pixr-train` console script, `python -m pixr.app.train`). """
+    code = _run(argv)
+    # NOTE only after a normal return, w/ an exception other ranks may still be in collectives, the launcher tears down
+    destroy_distributed()
+    return code
+
+
+if __name__ == '__main__':
+    sys.exit(main())

pixr/data/__init__.py

@@ -0,0 +1,2 @@
+from .loader import create_loader
+from .transforms import create_image_preprocessor, ImagePreprocessorFactory

pixr/data/loader.py

@@ -0,0 +1,26 @@
+from typing import Optional
+
+import chugr
+from chugr.common import DataArg, DataCfg, DataTaskCfg, DistributedCfg, LoaderBundle
+
+
+def create_loader(
+        data: DataArg,
+        task_cfg: DataTaskCfg,
+        is_training: bool = False,
+        start_interval: int = 0,
+        seed: int = 0,
+        world_size: int = 1,
+        global_rank: int = 0,
+        local_rank: int = 0,
+) -> LoaderBundle:
+    """ Create a chugr loader (webdataset or HF datasets) for the given task data pipeline config. """
+    data_cfg = DataCfg.from_arg(data)
+    return chugr.create_loader(
+        data_cfg,
+        task_cfg,
+        is_training=is_training,
+        start_interval=start_interval,
+        seed=seed,
+        distributed=DistributedCfg(world_size=world_size, local_rank=local_rank, global_rank=global_rank),
+    )

pixr/data/transforms.py

@@ -0,0 +1,52 @@
+from typing import Any, Dict, Optional
+
+import chugr
+from chugr.common import ImageInputCfg
+
+
+def create_image_preprocessor(
+        input_info: Dict[str, Any],
+        is_training: bool,
+        transform_type: Optional[str] = None,
+        max_seq_len: Optional[int] = None,
+):
+    """ Create image preprocessing for a model's image input requirements (see ImageEncoderTimm.input_info).
+
+    The same transform type is used for train and eval (eval disables augmentation), so input geometry
+    (resize, padding, patch sequence) always matches between training and evaluation.
+
+    Args:
+        input_info: model image input info
+        is_training: enable train augmentations
+        transform_type: chugr transform type, default 'doc_naflex' for NaFlex models, else 'doc_better'
+    """
+    naflex = input_info['naflex']
+    if transform_type is None:
+        transform_type = 'doc_naflex' if naflex else 'doc_better'
+    if naflex:
+        assert transform_type == 'doc_naflex', 'NaFlex models require doc_naflex transforms.'
+    cfg = ImageInputCfg(
+        size=input_info['image_size'],
+        mode=input_info['image_fmt'],
+        mean=input_info['mean'],
+        std=input_info['std'],
+        interpolation='bicubic',
+        fill_color=255,
+        transform_type=transform_type,
+        patch_size=input_info['patch_size'],
+        max_seq_len=max_seq_len or input_info['max_seq_len'],
+    )
+    return chugr.create_image_preprocessor(cfg, is_training=is_training)
+
+
+class ImagePreprocessorFactory:
+    """ seq_len -> image preprocessor, for NaFlex bucketed batching (picklable for dataloader workers). """
+
+    def __init__(self, input_info: Dict[str, Any], is_training: bool, transform_type: Optional[str] = None):
+        self.input_info = dict(input_info)
+        self.is_training = is_training
+        self.transform_type = transform_type
+
+    def __call__(self, seq_len: int):
+        return create_image_preprocessor(
+            self.input_info, self.is_training, transform_type=self.transform_type, max_seq_len=seq_len)

pixr/framework/__init__.py

@@ -0,0 +1,7 @@
+from .config import OptimizationCfg, TaskCfg, TrainerCfg
+from .device import DeviceEnv, DeviceEnvType, destroy_distributed
+from .logger import setup_logging
+from .monitor import Monitor
+from .random import random_seed
+from .task import TrainingTask, move_to_device
+from .trainer import Trainer

pixr/framework/config.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Tuple
+
+from pixr.models import ModelArgs
+from pixr.tokenizers import TokenizerCfg
+
+
+@dataclass
+class TaskCfg:
+    """ Base task config: model, tokenizer, preprocessing, and precision (used for training and eval). """
+    model: ModelArgs = field(default_factory=ModelArgs)
+    tokenizer: Optional[TokenizerCfg] = None  # default == text decoder name (or checkpoint tokenizer)
+    image_transforms: Optional[str] = None  # chugr transform type, default 'doc_naflex' or 'doc_better'
+    dtype: Optional[str] = 'bfloat16'  # low precision dtype for autocast ('bfloat16' or 'float16')
+    amp: bool = True
+    label_smoothing: float = 0.
+    z_loss_weight: float = 0.  # weight of the LM z-loss (mean logsumexp(logits) ** 2), e.g. 1e-4
+    # train text is trimmed to the longest sequence of the batch rounded up to this multiple, bounds the number of
+    # distinct text lengths (compile / FlexAttention mask shapes) for a few extra pad tokens, 1 disables rounding
+    text_pad_multiple: int = 64
+    loss_chunk_size: int = 4096  # fused (chunked) LM head + CE, tokens per chunk, 0 to materialize full logits
+
+
+@dataclass
+class OptimizationCfg:
+    optimizer: str = 'adamw'
+    scheduler: str = 'cosine'
+    learning_rate: float = 1e-4
+    warmup_learning_rate: float = 0.
+    min_learning_rate: float = 0.
+    weight_decay: float = .05
+    eps: float = 1e-8
+    clip_grad_value: Optional[float] = 1.0  # Gradient clipping value.
+    clip_grad_mode: Optional[str] = 'norm'  # Gradient clipping mode, one of ('norm', 'value', 'agc').
+    grad_accum_steps: int = 1
+    momentum: Optional[float] = None
+    betas: Optional[Tuple[float, float]] = (0.9, 0.98)
+    # per component LR multipliers (components defined by the model's param_layer_map(): 'encoder' == image
+    # encoder, 'bridge' == cross-attention / prefix projector, 'decoder' == text decoder / LM trunk)
+    encoder_lr_scale: float = 1.0
+    bridge_lr_scale: float = 1.0
+    decoder_lr_scale: float = 1.0
+    # layer-wise LR decay within the encoder / decoder (scale = decay ** (top layer - layer id)), None disables
+    encoder_layer_decay: Optional[float] = None
+    decoder_layer_decay: Optional[float] = None
+    # extra multipliers by param name regex (re.search, first match wins), e.g. {'embed_tokens': 2.0}
+    lr_scale_overrides: Optional[Dict[str, float]] = None
+    # scale loss by batch size relative to the mean for variable batch size (NaFlex schedule) loaders,
+    # 'none', 'linear' (per-sample gradient weight constant) or 'sqrt'
+    loss_batch_scale: str = 'none'
+
+
+@dataclass
+class TrainerCfg:
+    """ Optimization loop config, independent of the task. """
+    num_intervals: int = 100
+    num_warmup_intervals: float = 1.
+    log_frequency: int = 100  # log every n steps
+    metrics_frequency: int = 1000  # calculate (expensive) task train metrics every n steps, 0 disables
+    opt: OptimizationCfg = field(default_factory=OptimizationCfg)
+    grad_checkpointing: bool = False
+    ema_decay: Optional[float] = None  # enable EMA of model weights w/ this decay
+
+    # torch.compile, strategy (as per OpenCLIP --torchcompile-strategy):
+    #   'task': compile the task training forward (model + loss), after distributed wrapping
+    #   'model': compile the model before DDP wrapping
+    #   'blocks': compile blocks in place before DDP wrapping, grad checkpoint recompute stays in eager autograd
+    #     (bounds compiled backward memory when whole-graph compile schedules checkpointed recomputes poorly)
+    # w/ FSDP, blocks are always compiled per block inside prepare_fsdp() (compile -> composable AC -> FSDP)
+    compile: bool = False
+    compile_strategy: str = 'task'
+    compile_backend: Optional[str] = None
+    compile_mode: Optional[str] = None
+    compile_dynamic: Optional[bool] = None
+
+    # FSDP2 (fully_shard) instead of DDP for distributed training (params / grads / optimizer state sharded)
+    fsdp: bool = False
+    fsdp_reshard_after_forward: bool = True  # False keeps block params gathered from forward to backward
+    fsdp_offload_cpu: bool = False

pixr/framework/device.py

@@ -0,0 +1,176 @@
+""" Distributed training/validation utils
+
+"""
+import os
+from dataclasses import dataclass, field, InitVar
+from enum import Enum
+from typing import Union, Optional, List, Tuple
+
+import torch
+import torch.distributed as dist
+
+
+def is_distributed_env():
+    if "WORLD_SIZE" in os.environ:
+        return int(os.environ["WORLD_SIZE"]) > 1
+    if "SLURM_NTASKS" in os.environ:
+        return int(os.environ["SLURM_NTASKS"]) > 1
+    return False
+
+
+def world_info_from_env():
+    local_rank = 0
+    for v in (
+        "LOCAL_RANK",
+        "MPI_LOCALRANKID",
+        "SLURM_LOCALID",
+        "OMPI_COMM_WORLD_LOCAL_RANK",
+    ):
+        if v in os.environ:
+            local_rank = int(os.environ[v])
+            break
+
+    global_rank = 0
+    for v in ("RANK", "PMI_RANK", "SLURM_PROCID", "OMPI_COMM_WORLD_RANK"):
+        if v in os.environ:
+            global_rank = int(os.environ[v])
+            break
+
+    world_size = 1
+    for v in ("WORLD_SIZE", "PMI_SIZE", "SLURM_NTASKS", "OMPI_COMM_WORLD_SIZE"):
+        if v in os.environ:
+            world_size = int(os.environ[v])
+            break
+
+    return local_rank, global_rank, world_size
+
+
+class DeviceEnvType(Enum):
+    """Device Environment Types"""
+
+    CPU = "cpu"
+    CUDA = "cuda"
+    XLA = "xla"
+
+
+@dataclass
+class DeviceEnv:
+    """
+    Represents the environment for device setup and distributed processing.
+
+    Attributes:
+        device (torch.device): Current device (e.g., CUDA GPU) that the model/data should reside in.
+        world_size (Optional[int]): Total number of processes involved in the distributed setup.
+        local_rank (Optional[int]): Rank of the current process on the current node.
+        global_rank (Optional[int]): Rank of the current process globally across all nodes.
+
+    Init Arguments:
+        init_device_type (Optional[str]): The type of device to initialize (e.g., 'cuda', 'cpu').
+            Defaults to None, which means it will be inferred.
+        init_device_index (Optional[int]): Index of the CUDA device to initialize if multiple are present.
+            Defaults to None, indicating the first available device.
+        init_dist_backend (str): Backend for distributed training. Default is 'nccl'.
+        init_dist_url (str): URL for initializing distributed training. Default is 'env://'.
+
+    Methods:
+        is_global_primary: Checks if current process has the global rank 0.
+        is_local_primary: Checks if current process has the local rank 0.
+        is_primary: Checks if current process is either a local or global primary.
+        broadcast_object: Broadcasts an object from a source process to all other processes.
+        all_gather_object: Gathers objects from all processes and returns them.
+    """
+
+    init_device_type: InitVar[Optional[str]] = None
+    init_device_index: InitVar[Optional[int]] = None
+    init_dist_backend: InitVar[str] = "nccl"
+    init_dist_url: InitVar[str] = "env://"
+
+    device: torch.device = field(
+        init=False
+    )  # set from device_type + device_index or post_init logic
+    world_size: Optional[int] = None  # set by post_init from env when None
+    local_rank: Optional[int] = None  # set by post_init from env when None
+    global_rank: Optional[int] = None  # set by post_init from env when None
+
+    def is_global_primary(self):
+        return self.global_rank == 0
+
+    def is_local_primary(self):
+        return self.local_rank == 0
+
+    def is_primary(self, local=False):
+        return self.is_local_primary() if local else self.is_global_primary()
+
+    def __post_init__(
+        self,
+        init_device_type: Optional[str],
+        init_device_index: Optional[int],
+        init_dist_backend: str,
+        init_dist_url: str,
+    ):
+        # FIXME support different device types, just using cuda to start
+        assert torch.cuda.device_count()
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cuda.matmul.allow_tf32 = True
+        init_local_rank, init_global_rank, init_world_size = world_info_from_env()
+        if init_world_size > 1:
+            # setup distributed
+            assert init_device_index is None
+            self.local_rank = int(init_local_rank)
+            is_slurm = "SLURM_PROCID" in os.environ
+            if "SLURM_PROCID" in os.environ:
+                # os.environ['LOCAL_RANK'] = str(init_local_rank)
+                # os.environ['RANK'] = str(init_global_rank)
+                # os.environ['WORLD_SIZE'] = str(init_world_size)
+                torch.distributed.init_process_group(
+                    backend=init_dist_backend,
+                    init_method=init_dist_url,
+                    world_size=init_world_size,
+                    rank=init_global_rank,
+                )
+            else:
+                torch.distributed.init_process_group(
+                    backend=init_dist_backend,
+                    init_method=init_dist_url,
+                )
+
+            self.world_size = torch.distributed.get_world_size()
+            self.global_rank = torch.distributed.get_rank()
+            if is_slurm:
+                assert self.world_size == init_world_size
+                assert self.global_rank == init_global_rank
+
+            self.device = torch.device("cuda:%d" % self.local_rank)
+            torch.cuda.set_device(self.local_rank)
+        else:
+            self.device = torch.device(
+                "cuda" if init_device_index is None else f"cuda:{init_device_index}"
+            )
+            self.local_rank = 0
+            self.world_size = 1
+            self.global_rank = 0
+
+    def broadcast_object(self, obj, src=0):
+        # broadcast a pickle-able python object from rank-0 to all ranks
+        if self.global_rank == src:
+            objects = [obj]
+        else:
+            objects = [None]
+        dist.broadcast_object_list(objects, src=src)
+        return objects[0]
+
+    def all_gather_object(self, obj, dst=0):
+        # gather a pickle-able python object across all ranks
+        objects = [None for _ in range(self.world_size)]
+        dist.all_gather_object(objects, obj)
+        return objects
+
+
+def destroy_distributed():
+    """ Tear down the default process group (if any) before interpreter exit.
+
+    W/o this, the NCCL watchdog thread can still be querying events while CUDA unloads at exit, crashing a rank
+    ('driver shutting down', SIGSEGV) after a successful run.
+    """
+    if dist.is_available() and dist.is_initialized():
+        dist.destroy_process_group()