Snowflake Container Runtime GPU Version 2.6 (Latest)¶

Added¶

• New Ray Train-based PyTorch distributed training backend in preview, opt-in via env variable PYTORCH_USE_LEGACY_TRAINER=0. The legacy custom trainer remains the default (PYTORCH_USE_LEGACY_TRAINER=1) for this release; the default will flip to the Ray Train backend in an upcoming release.
• train_func may now optionally accept a config dict argument (Ray Train convention); zero-argument functions continue to work.

Changed¶

• BREAKING: Opting into the Ray Train PyTorch backend (PYTORCH_USE_LEGACY_TRAINER=0) introduces these differences from the legacy trainer:
  • Ray Train calls init_process_group() automatically — guard user calls with if not torch.distributed.is_initialized() and remove/guard destroy_process_group().
  • trainer.run(use_ray_job=...) is ignored; training always runs in the driver process.
  • ray.train.torch.prepare_data_loader() auto-moves batches to GPU — pass move_to_device=False or create all intermediate tensors on the target device to avoid device-mismatch crashes.

Deprecated¶

• Legacy custom PyTorch trainer (PYTORCH_USE_LEGACY_TRAINER=1). Still the current default, but planned for removal once the Ray Train backend becomes the default in an upcoming release. Migrate training functions per the breaking-change notes above.
• get_context().get_metrics_reporter() — unavailable on the Ray Train backend. Use ray.train.report({"loss": ...}) instead.
• get_context().get_rank() — use get_world_rank().
• get_context().get_master_addr() / get_master_port() — read os.environ["MASTER_ADDR"] / os.environ["MASTER_PORT"] directly.
• get_context().get_default_backend() / get_supported_backends() — use torch.distributed.get_backend() after process group initialization.

The following lists the packages available for each Python version of GPU version 2.6.