Fine-Tuning LLMs on Snowflake with QLoRA and FSDP2

Most teams don’t need to fine-tune a Large Language Model (LLM). Most managed LLM APIs cover the vast majority of use cases at a fraction of the operational cost. But sometimes you do need a custom adapter: e.g., a domain-specific vocabulary, a proprietary dataset that can’t leave your environment, or a task narrow enough that a small fine-tuned model beats a general-purpose one on latency and cost.

This post walks through SFT Trainer, a small Python library built to run QLoRA supervised fine-tuning (SFT) jobs directly on Snowflake, using SPCS (Snowpark Container Services) GPU nodes and FSDP2 for distributed training. It’s less about the library itself and more about two things that tend to get skipped in “fine-tune your own LLM” tutorials: how you actually profile a multi-GPU job to find out what’s slow, and how the profiler’s findings should shape the GPU environment you provision in the first place.

Why train inside Snowflake at all?

The obvious alternative is to export training data, spin up GPUs somewhere else (SageMaker, a GPU cloud, your own cluster), train, then import the resulting weights back. That works, but it means:

• Data leaves the governed environment (or you build a separate pipeline to keep it in sync).
• You maintain two infrastructures: the warehouse/data platform, and the training platform.
• Checkpoints, experiment metadata and traces live somewhere your data team can’t query with SQL.

Running training as an SPCS job keeps everything (data, compute, checkpoints, metrics) inside the same governance boundary. The trade-off is that you have to bring your own training loop into a container-based compute pool instead of using a managed training service. SFT Trainer is that training loop.

flowchart LR
    subgraph Snowflake["Snowflake account"]
        TT["Training data table\n(PROMPT / COMPLETION)"]
        ST["Internal stage\n(checkpoints, traces)"]
        ET["Experiment tables\n(metrics, params)"]
        subgraph SPCS["SPCS GPU compute pool"]
            J["SFT Trainer job\n(Lightning + FSDP2)"]
        end
    end

    TT -- "read via Snowpark session" --> J
    J -- "upload LoRA checkpoints" --> ST
    ST -- "resume / fetch_to_local" --> J
    J -- "log metrics & params" --> ET
    J -- "upload profiler traces" --> ST

# Run via the CLI
python -m sft_trainer train-sft --log-level DEBUG

All settings are configured through prefixed environment variables, with nested settings using a single underscore delimiter (e.g. LLM_MODEL_ID, DATA_BATCH_SIZE, LOGGER_BACKEND). This makes the same training code portable across notebooks, CLI jobs, and SPCS services. You just change environment variables, not code.

QLoRA: fine-tuning without the GPU memory bill

Full fine-tuning of an 8B+ parameter model requires storing optimizer states, gradients, and weights in full precision; hundreds of GB of GPU memory. QLoRA sidesteps this:

1. The base model is loaded in 4-bit quantized form and frozen.
2. Small trainable low-rank adapter matrices (LoRA) are injected into attention/MLP layers.
3. Only the LoRA adapters are updated during backpropagation.

flowchart TB
    W["Frozen base weights\n(4-bit quantized)"] --> F["Forward pass"]
    A["LoRA adapters A, B\n(fp16/bf16, low-rank)"] --> F
    F --> L["Loss"]
    L --> B["Backward pass"]
    B -.->|"gradients flow only here"| A
    B -.->|"no gradient"| W

The practical consequence: checkpoints only need to persist the adapter weights, not the full model. Here’s the actual LightningModule, trimmed to the parts that matter:

class SFTQLoRAModule(SnowflakeSessionModel):
    """QLoRA supervised fine-tuning module using torchao quantization and FSDP2."""

    def configure_model(self) -> None:
        """Load, quantize, and wrap the base model with LoRA adapters.

        In PyTorch Lightning 2.x, configure_model() is the hook where FSDP2
        intercepts and shards the model, so model construction must happen
        here rather than in __init__.
        """
        base_model = AutoModelForCausalLM.from_pretrained(
            self.model_id,
            torch_dtype=torch.bfloat16,
            device_map=None,
        )

        if self._quantize:
            quantize_(base_model, Int4WeightOnlyConfig())

        lora_config = LoraConfig(
            r=self.lora_r,
            lora_alpha=self.lora_alpha,
            target_modules=self.lora_target_modules,
            task_type=TaskType.CAUSAL_LM,
            lora_dropout=self.lora_dropout,
        )
        self.model = get_peft_model(base_model, lora_config)

    def training_step(self, batch: dict[str, torch.Tensor], batch_idx: int) -> STEP_OUTPUT:
        """Run a single training step and log the loss."""
        outputs = self.model(
            input_ids=batch["input_ids"], attention_mask=batch["attention_mask"], labels=batch["labels"]
        )
        loss: torch.Tensor = outputs.loss
        self.log("train/loss", loss, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True)
        return loss

    def state_dict(self) -> dict[str, torch.Tensor]:
        """Return only LoRA adapter weights, keeping checkpoints small."""
        full = super().state_dict()
        return {k: v for k, v in full.items() if "lora_" in k}

    def load_state_dict(self, state_dict: dict[str, torch.Tensor], strict: bool = False) -> object:
        """Load LoRA adapter weights onto the base model (rebuilt in configure_model)."""
        return super().load_state_dict(state_dict, strict=False)

Two details worth calling out:

• Model construction happens in configure_model(), not __init__(). That’s not a style choice, it’s a Lightning 2.x requirement: FSDP2 needs to intercept the model after it’s built but before training starts, so it can shard it across ranks. Building it eagerly in __init__ would sidestep that hook entirely.
• state_dict() filters on the "lora_" substring in parameter names, which is all peft’s get_peft_model() needs to be able to reconstruct the adapter later. This is what shrinks checkpoints from gigabytes to kilobytes: the frozen 4-bit base weights simply never get written to disk.

This single design choice is what makes checkpoint upload/download to a Snowflake stage on every save cheap enough to do routinely, instead of only at the end of training.

Checkpointing to a Snowflake stage

SFT Trainer plugs a custom SnowflakeStageCheckpointIO into PyTorch Lightning’s Trainer, so checkpoints transparently upload to (and download from) an internal stage:

from sft_trainer.model import SFTQLoRAModule
from sft_trainer.plugins import SnowflakeStageCheckpointIO

checkpoint_io = SnowflakeStageCheckpointIO(session=session, stage="@DB.SCHEMA.STAGE/checkpoints")

trainer = L.Trainer(plugins=[checkpoint_io])
trainer.fit(module, train_dataloaders=dataloader)

Underneath, SnowflakeStageCheckpointIO is a small subclass of Lightning’s CheckpointIO interface: three methods are all it takes:

class SnowflakeStageCheckpointIO(CheckpointIO):
    """Saves and loads Lightning checkpoints to/from a Snowflake stage."""

    def __init__(self, session: Session, stage: str) -> None:
        self._session = session
        self._stage = stage

    def save_checkpoint(self, checkpoint: dict[str, Any], path: str | Path, storage_options=None) -> None:
        """Save checkpoint locally and upload to Snowflake stage."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(checkpoint, path)

        logger.info(f"Uploading checkpoint to stage: {self._stage}/{path.name}")
        self._session.file.put(f"file://{path}", self._stage, auto_compress=False, overwrite=True)

    def load_checkpoint(self, path: str | Path, map_location=None, weights_only=None) -> dict[str, Any]:
        """Load checkpoint from stage with distributed-aware downloading.

        Only rank 0 downloads from stage; the result is broadcast to all
        other ranks so N GPUs don't all try to hit the stage at once.
        """
        is_distributed = dist.is_available() and dist.is_initialized()
        rank = dist.get_rank() if is_distributed else 0
        checkpoint: dict[str, Any] | None = None

        if rank == 0:
            checkpoint = self._download_and_load(path, map_location, weights_only or False)

        if is_distributed:
            object_list = [checkpoint]
            dist.broadcast_object_list(object_list, src=0)
            checkpoint = object_list[0]

        assert checkpoint is not None
        return checkpoint

    def fetch_to_local(self, path: str | Path) -> Path:
        """Download a checkpoint from stage to the local path (FSDP resume path).

        FSDP bypasses CheckpointIO for loading and reads directly from the
        filesystem, so this must be called explicitly before Trainer.fit
        when resuming under FSDP.
        """
        path = Path(path)
        is_distributed = dist.is_available() and dist.is_initialized()
        rank = dist.get_rank() if is_distributed else 0

        if rank == 0:
            stage_source = f"{self._stage}/{path.name}"
            path.parent.mkdir(parents=True, exist_ok=True)
            self._session.file.get(stage_source, f"file://{path.parent}")

        if is_distributed:
            dist.barrier()  # all ranks wait for rank 0's download

        return path

The dist.barrier() call at the end of fetch_to_local is easy to miss but is what actually makes multi-node resume correct: without it, ranks 1..N could start reading a checkpoint file from the local filesystem before rank 0 has finished writing it.

The separation of concerns is deliberate:

Concern	Who handles it
What to save (LoRA-only state)	SFTQLoRAModule.state_dict()
When to save	Lightning’s ModelCheckpoint callback
Where to persist	SnowflakeStageCheckpointIO

For multi-GPU/multi-node runs, only rank 0 uploads on save. On load, DDP broadcasts the checkpoint dict to all ranks via torch.distributed; FSDP requires an explicit fetch_to_local() call before Trainer.fit, since FSDP bypasses the CheckpointIO abstraction on load.

sequenceDiagram
    participant R0 as Rank 0
    participant Rn as Rank 1..N
    participant Stage as Snowflake stage

    Note over R0,Rn: Save (every N epochs)
    R0->>Stage: upload checkpoint (LoRA weights only)
    Rn->>Rn: no-op (rank 0 owns upload)

    Note over R0,Rn: Resume (DDP)
    R0->>Stage: download checkpoint
    R0->>Rn: broadcast checkpoint dict (torch.distributed)

    Note over R0,Rn: Resume (FSDP)
    R0->>Stage: fetch_to_local()
    R0->>Rn: barrier sync
    Rn->>Rn: read checkpoint from local filesystem

Diagnosing distributed training with a Chrome trace profiler

Multi-GPU training introduces a whole new failure mode: the model trains, but slowly, and it’s not obvious why. SnowflakeProfiler extends Lightning’s PyTorchProfiler to capture Chrome-trace JSON files (every CUDA kernel, every NCCL collective) plus diagnostic dumps (nvidia-smi, ibstat output) at teardown, then uploads them to a stage automatically.

What the profiler actually captures

A training step is instrumented in three phases, controlled independently so the trace stays small enough to actually open:

flowchart LR
    S["Step 0..N"] --> Wait["wait steps\n(skip, let JIT warm up)"]
    Wait --> Warm["warmup steps\n(profiler attached, data discarded)"]
    Warm --> Active["active steps\n(fully recorded)"]
    Active --> Repeat{"repeat cycle?"}
    Repeat -- yes --> Wait
    Repeat -- no --> Done["upload trace to stage"]

Each active step records, per GPU and per CPU thread:

• Every aten::* op dispatched (matmuls, layernorms, attention kernels).
• Every CUDA kernel launched and its wall-clock duration on the device stream.
• Every NCCL collective (AllReduce, AllGather, ReduceScatter) and which ranks participated.
• Host-to-device (HtoD) and device-to-host (DtoH) memory copies.

The scheduling itself is just standard torch.profiler.schedule, wired up in the training entrypoint:

profiler = SnowflakeProfiler(
    dirpath=settings.profiler.trace_dir,
    filename=settings.profiler.trace_filename,
    schedule=torch.profiler.schedule(
        wait=settings.profiler.wait_steps,
        warmup=settings.profiler.warmup_steps,
        active=settings.profiler.active_steps,
        repeat=settings.profiler.repeat,
    ),
    export_to_chrome=True,
    record_shapes=True,
    profile_memory=True,
)

SnowflakeProfiler itself is a thin subclass of Lightning’s PyTorchProfiler. All it adds is running a batch of shell diagnostics and writing their stdout next to the Chrome trace when the profiler tears down:

class SnowflakeProfiler(PyTorchProfiler):
    """PyTorchProfiler that also captures GPU/network diagnostics on teardown."""

    def __init__(self, *args, extra_cmds: list[DiagnosticCommand] | None = None, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.extra_cmds = extra_cmds or settings.profiler.diagnostics

    def _capture_diagnostics(self) -> None:
        """Run diagnostic commands and write stdout to dirpath."""
        if self.dirpath is None:
            return
        trace_path = Path(self.dirpath)
        trace_path.mkdir(parents=True, exist_ok=True)

        for diag in self.extra_cmds:
            try:
                result = subprocess.run(diag.command, capture_output=True, text=True, timeout=30)
                if result.returncode == 0:
                    (trace_path / diag.filename).write_text(result.stdout)
                else:
                    logger.debug(f"Skipping {diag.filename}: exited with code {result.returncode}")
            except (FileNotFoundError, subprocess.TimeoutExpired) as e:
                logger.debug(f"Skipping {diag.filename}: {e}")

    def teardown(self, stage: str | None = None) -> None:
        """Capture GPU diagnostics then tear down the profiler."""
        self._capture_diagnostics()
        super().teardown(stage=stage)

Note the try/except around each command: ibstat doesn’t exist on nodes without InfiniBand hardware, and that’s expected on single-node runs, not an error. The diagnostics themselves are declared as data, not hardcoded shell strings, which is what lets them be overridden per environment:

diagnostics: list[DiagnosticCommand] = Field(
    default_factory=lambda: [
        DiagnosticCommand(filename="gpu_topology.txt", command=["nvidia-smi", "topo", "-m"]),
        DiagnosticCommand(filename="nvlink_status.txt", command=["nvidia-smi", "nvlink", "-s"]),
        DiagnosticCommand(
            filename="gpu_clocks_throttle.txt", command=["nvidia-smi", "-q", "-d", "PERFORMANCE"]
        ),
        DiagnosticCommand(
            filename="gpu_utilization.txt", command=["nvidia-smi", "dmon", "-s", "u", "-c", "5"]
        ),
        DiagnosticCommand(filename="ibstat.txt", command=["ibstat"]),
    ],
)

Five files, five questions, every single run: is NVLink actually being used (gpu_topology.txt), is it saturated (nvlink_status.txt), is any GPU throttled (gpu_clocks_throttle.txt), is utilization balanced across ranks (gpu_utilization.txt), and is InfiniBand even up (ibstat.txt)? These land on the stage automatically; there’s no need to SSH into a node mid-run to remember to capture them.

Opening the Chrome trace itself in chrome://tracing or Perfetto turns an abstract “training is slow” complaint into a concrete picture:

CUDA stream 7:  |==kernel==|                              |==kernel==|
                           ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
                           GPU idle: CPU can't feed work fast enough

From symptom to root cause: a triage flow

Rather than reading traces ad hoc, it helps to walk a fixed decision tree every time a run looks slow. This is roughly what happens, in order, whenever step time in the experiment logs jumps unexpectedly:

flowchart TD
    Start["Step time regressed"] --> Q1{"Gaps on\nCUDA stream?"}
    Q1 -- yes --> Q2{"Gap precedes\nforward pass?"}
    Q2 -- yes --> DataStarve["Data-starved\nDataLoader too slow"]
    Q2 -- no --> Q3{"Gap coincides with\nNCCL bar?"}
    Q3 -- yes --> CommBound["Communication-bound\nAllReduce/AllGather dominates"]
    Q3 -- no --> KernelLaunch["Launch-overhead bound\nmany tiny kernels"]

    Q1 -- no --> Q4{"HtoD copies\ninside step?"}
    Q4 -- yes --> HostCopy["CPU-resident tensors\nmoved every step"]
    Q4 -- no --> Q5{"One rank\nfinishes early,\nwaits at barrier?"}
    Q5 -- yes --> Straggler["Straggler GPU\nthermal or batch imbalance"]
    Q5 -- no --> Topology["Check nvidia-smi topo -m\nPCIe vs NVLink path"]

    DataStarve --> FixData["num_workers, pin_memory,\nprefetch_factor"]
    CommBound --> FixComm["NCCL_IB_HCA, HYBRID_SHARD,\ngradient accumulation"]
    KernelLaunch --> FixKernel["torch.compile,\nCUDA graphs"]
    HostCopy --> FixHost["precompute on-device,\nregister_buffer"]
    Straggler --> FixStrag["nvidia-smi -q -d PERFORMANCE\ncheck throttling"]
    Topology --> FixTopo["gradient accumulation,\nHYBRID_SHARD"]

A few patterns come up repeatedly, matched to that tree:

• GPU idle gaps before the forward pass → the DataLoader can’t keep up. Fix: more workers, pin_memory=True, higher prefetch_factor.
• Wide NCCL AllReduce bars relative to compute → communication-bound, usually a cross-node bottleneck. Check NCCL_IB_HCA is pointing at the right InfiniBand adapter, or switch FSDP2 to HYBRID_SHARD so sharding stays within a node (fast NVLink) and only gradients are reduced across nodes.
• Hundreds of tiny kernels with visible launch gaps → CPU-side cudaLaunchKernel overhead dominates. Fuse ops with torch.compile(mode="reduce-overhead") or enable CUDA graphs.
• HtoD copies mid-step, not just during data loading → tensors (masks, positional encodings) are being created on CPU and moved every step. Precompute once and cache on-device, or use register_buffer.
• Straggler ranks waiting at the AllReduce barrier → uneven batch/sequence lengths or thermal throttling; nvidia-smi -q -d PERFORMANCE shows the throttle reason directly.
• nvidia-smi topo -m shows PHB/SYS instead of NV# → the GPU pair is on PCIe, not NVLink, roughly a 20x bandwidth difference. Gradient accumulation amortizes the slower link over more compute per sync.

None of this is Snowflake-specific; it’s standard distributed-training debugging. But having the traces land automatically on a stage next to the training data means you don’t need a separate observability stack to look at them.

From profiler findings to GPU environment architecture

This is the part that’s easy to skip: profiling isn’t just for fixing a slow run in place, it’s the input to deciding what to provision next time. The trace tells you which resource is the bottleneck (compute, host-to-device bandwidth, intra-node interconnect, or inter-node network), and each of those points to a different fix on the infrastructure side, not just the code side.

flowchart TD
    Trace["Profiler trace + nvidia-smi/ibstat dumps"] --> Bottleneck{"Dominant\nbottleneck?"}

    Bottleneck -- "GPU idle,\nCPU-bound" --> Env1["More vCPU per GPU\nor separate CPU-heavy preprocessing node"]
    Bottleneck -- "Intra-node\nNVLink saturated" --> Env2["Model fits in fewer,\nlarger-memory GPUs per node"]
    Bottleneck -- "Cross-node\nAllReduce dominates" --> Env3["Fewer, bigger nodes\nor HYBRID_SHARD topology"]
    Bottleneck -- "PCIe fallback\n(no NVLink)" --> Env4["Choose instance family\nwith NVLink-connected GPUs"]
    Bottleneck -- "Compute-bound,\nGPUs saturated" --> Env5["Scale out: add nodes\n(compute genuinely needed)"]

    Env1 --> Decision["Update compute pool spec"]
    Env2 --> Decision
    Env3 --> Decision
    Env4 --> Decision
    Env5 --> Decision

In practice, this maps onto a handful of concrete topology choices for the SPCS compute pool:

Single node, NVLink-connected GPUs

If the trace shows communication bars that are already fast (NVLink, not PCIe) and the bottleneck is elsewhere, a single multi-GPU node is enough; there’s no need to pay for cross-node InfiniBand at all.

flowchart TB
    subgraph Node["Single SPCS node (GPU_NV_M)"]
        CPU["CPU / host memory"]
        G0["GPU 0"]
        G1["GPU 1"]
        G2["GPU 2"]
        G3["GPU 3"]
        CPU <--> G0
        CPU <--> G1
        G0 <-->|"NVLink"| G1
        G1 <-->|"NVLink"| G2
        G2 <-->|"NVLink"| G3
        G3 <-->|"NVLink"| G0
    end

Multi-node with InfiniBand, FULL_SHARD

When the model doesn’t fit in one node’s aggregate memory, FSDP2’s FULL_SHARD shards parameters across every rank, everywhere, which means every AllReduce/AllGather crosses the InfiniBand fabric, not just NVLink. This is correct when you genuinely need the extra memory, but it’s the most network-hungry topology available.

flowchart TB
    subgraph N1["Node 1"]
        A0["GPU 0"] <-->|NVLink| A1["GPU 1"]
    end
    subgraph N2["Node 2"]
        B0["GPU 0"] <-->|NVLink| B1["GPU 1"]
    end
    A0 <-->|"InfiniBand\n(full param shards)"| B0
    A1 <-->|"InfiniBand\n(full param shards)"| B1

Multi-node with InfiniBand, HYBRID_SHARD

If the trace shows the InfiniBand link is the bottleneck (wide AllReduce bars, ibstat shows a healthy link but it’s still saturated), HYBRID_SHARD shards fully within a node (fast NVLink) and only replicates/reduces gradients across nodes. Less cross-node traffic, at the cost of higher per-node memory since parameters are duplicated rather than fully sharded.

flowchart TB
    subgraph N1["Node 1 (full model shard)"]
        A0["GPU 0"] <-->|"NVLink\n(param shards)"| A1["GPU 1"]
    end
    subgraph N2["Node 2 (replica of shard)"]
        B0["GPU 0"] <-->|"NVLink\n(param shards)"| B1["GPU 1"]
    end
    A0 -.->|"InfiniBand\n(gradients only)"| B0

The rule of thumb that comes out of repeated profiling sessions: default to a single, NVLink-dense node until the trace actually shows the model is memory-constrained across nodes. Only then pay for InfiniBand bandwidth, and prefer HYBRID_SHARD over FULL_SHARD unless per-node memory genuinely can’t hold a full shard replica.

Turning a profiler finding into a one-line config change

The whole point of separating “trace analysis” from “infrastructure change” is that the fix should never require touching training code. Here’s the actual settings model backing the sharding strategy, with the trade-off spelled out in the field description itself so nobody has to go dig through FSDP docs mid-incident:

class TrainerSettings(BaseModel):
    """Lightning Trainer and distributed strategy configuration."""

    strategy: StrategyType = Field("fsdp2", description="Lightning distributed strategy to use.")
    fsdp_sharding_strategy: str = Field(
        "FULL_SHARD",
        description=(
            "FSDP sharding strategy. Options: 'FULL_SHARD' (shard params+grads+optim "
            "across all ranks), 'SHARD_GRAD_OP' (shard grads+optim only), 'HYBRID_SHARD' "
            "(full-shard intra-node via NVLink, replicate across nodes to reduce "
            "InfiniBand traffic), 'NO_SHARD' (DDP-like, no sharding)."
        ),
    )
    accumulate_grad_batches: int = Field(
        1,
        description=(
            "Number of batches to accumulate gradients over before an optimizer step. "
            "Effectively multiplies batch size without increasing memory usage, and "
            "amortizes AllReduce cost over more compute (useful for PCIe-bound setups)."
        ),
    )

And here’s where that setting actually becomes a live FSDP strategy object, straight from the training entrypoint:

resolved_strategy: str | Strategy
if settings.trainer.strategy in ("fsdp", "fsdp2"):
    resolved_strategy = FSDPStrategy(
        sharding_strategy=settings.trainer.fsdp_sharding_strategy,
    )
else:
    resolved_strategy = settings.trainer.strategy

trainer = L.Trainer(
    accelerator=settings.trainer.accelerator,
    devices=settings.trainer.devices,
    strategy=resolved_strategy,
    precision=settings.trainer.precision,
    accumulate_grad_batches=settings.trainer.accumulate_grad_batches,
    profiler=profiler,
    plugins=[checkpoint_io],
)

Which means the fix for a communication-bound trace really is a one-line environment variable change on redeploy, not a code review:

# Before: default FULL_SHARD, InfiniBand saturated on every AllGather
TRAINER_FSDP_SHARDING_STRATEGY=FULL_SHARD python -m sft_trainer train-sft

# After: switch to HYBRID_SHARD once the trace confirms cross-node traffic is the bottleneck
TRAINER_FSDP_SHARDING_STRATEGY=HYBRID_SHARD python -m sft_trainer train-sft

The same pattern applies to the NCCL transport settings that ibstat/NCCL_IB_HCA diagnostics feed into. These aren’t scattered os.environ[...] calls sprinkled through the codebase; they’re declared once, validated once, and exported once on startup:

class NcclSettings(BaseModel):
    """NCCL debugging and transport configuration for distributed training diagnostics."""

    debug: str = Field("WARN", description="NCCL log verbosity. Set to 'INFO' to diagnose AllReduce hangs.")
    debug_subsys: str = Field("ALL", description="NCCL subsystems to log (INIT, COLL, P2P, ALL, etc.).")
    ib_hca: str | None = Field(
        None,
        description=(
            "InfiniBand Host Channel Adapter(s) NCCL should use for cross-node AllReduce. "
            "Examples: 'mlx5_0', 'mlx5_0,mlx5_1' (round-robin), '^mlx5_0:2' (exclude port 2)."
        ),
    )


class Settings(BaseSettings):
    nccl: NcclSettings = Field(default_factory=NcclSettings)

    def model_post_init(self, __context: object) -> None:
        """Export environment variables that third-party libraries read at runtime."""
        super().model_post_init(__context)
        os.environ["NCCL_DEBUG"] = self.nccl.debug
        os.environ["NCCL_DEBUG_SUBSYS"] = self.nccl.debug_subsys
        if self.nccl.ib_hca is not None:
            os.environ["NCCL_IB_HCA"] = self.nccl.ib_hca

So the moment ibstat on a new node reports mlx5_1 instead of mlx5_0, the fix is NCCL_IB_HCA=mlx5_1 in the pool’s environment, not a patch to the training script.

Concrete diagnostics behind each decision

# Confirm intra-node link type before assuming NVLink is available
nvidia-smi topo -m
# NV# = NVLink (fast), PHB/SYS = PCIe (20x slower)

# Confirm InfiniBand is active before blaming NCCL config
ibstat
# Look for: State: Active, Physical state: LinkUp, Rate: 200 (Gb/s)

# Confirm NCCL is actually using the InfiniBand HCA, not falling back to TCP
echo $NCCL_SOCKET_IFNAME   # should match the IB interface (e.g. ib0)
echo $NCCL_IB_HCA          # should list the HCA (e.g. mlx5_0)

# Confirm no single rank is thermal-throttled before blaming batch imbalance
nvidia-smi -q -d PERFORMANCE

If ibstat shows no active ports, InfiniBand simply isn’t available on that node family, and NCCL silently falls back to TCP over Ethernet: orders of magnitude slower for AllReduce. That’s an infrastructure decision (pick a node family with InfiniBand), not something HYBRID_SHARD or gradient accumulation can fix.

Experiment tracking without leaving SQL

The default logger backend writes metrics and params straight to Snowflake tables via snowflake.ml.experiment.ExperimentTracking, so a hyperparameter sweep’s results are just… queryable. An MLflow backend is also available for teams standardized on it, syncing to a local ./mlruns directory and periodically pushing to a stage, at the cost of no longer being queryable via SQL.

For sweeps themselves, hydra-zen composes configs programmatically instead of maintaining a YAML tree per run:

from hydra_zen import builds, instantiate
from sft_trainer.config import Settings

settings_conf = builds(Settings, populate_full_signature=True)

for lr in [1e-4, 2e-4, 5e-4]:
    settings = instantiate(settings_conf, learning_rate=lr)
    # launch SPCS job with settings...

Because every run’s profiler trace, checkpoints, and metrics land in the same account, a sweep across GPU topologies (single-node NVLink vs. multi-node HYBRID_SHARD vs. multi-node FULL_SHARD) can be compared the same way as a sweep across learning rates: by querying the experiment tables, not by manually collecting logs from three different clusters.

Is it worth it? A cost check

Self-hosting only makes sense above a certain usage threshold. Based on actual account usage, the rough numbers looked like this:

Approach	Monthly Credits
Cortex AI (pay-per-token)	~12
Self-hosted 8B quantized, 24/7	~1,080
Self-hosted 8B quantized, 8hr/weekday	~264

At low-to-moderate usage, a managed LLM API wins by 20-90x. The break-even point was roughly 500 credits/month of sustained Cortex AI spend, around 40x the usage observed here. Self-hosting (and everything above) only becomes worth building when at least one of these holds:

• Batch inference over >100k rows/day.
• Latency-sensitive serving requiring sub-100ms p99 (avoiding managed-API cold starts).
• A fine-tuned adapter that isn’t available through the managed service at all.
• Data residency constraints requiring inference to stay inside a private network boundary.

Takeaway

The interesting part of this project isn’t QLoRA or FSDP2; both are well-documented elsewhere. It’s that none of the surrounding infrastructure (checkpoints, profiler traces, experiment metrics, training data) had to leave the data platform to make distributed GPU training work, and that the same profiler traces that fix today’s slow run also tell you what to provision for the next one. That’s a reasonable trade to make once you’ve confirmed self-hosting is actually justified, and not a moment before.