AWS Machine Learning Specialty (AWS): Tuning 60 XGBoost Configurations With a Six-Hour Budget

Where the time goes now: 45 min/trial x 60 = 45 hours serial. Even at 8-way parallelism that is about 5.6 hours with zero slack and no final retrain. Attack the per-trial cost first.

1. Kill the download. Switch to Pipe mode or sharded protobuf so training streams from S3; the 12-minute prefix disappears and the trial drops to roughly 30 minutes.
2. Tune on a stratified 15-20% subsample. That brings a trial to about 6-10 minutes. Sixty trials at 8-10 concurrent (max_parallel_jobs) with Bayesian optimization finishes well inside the budget.
3. Bound every trial with max_run and enable early stopping on the validation metric so bad configurations die early instead of running to completion.
4. Warm start from last quarter's tuning job; reuse its top configurations rather than re-exploring known-bad regions.
5. Reserve the last 1.5 hours to retrain the winning configuration on 100% of the data.

Transfer caveat: regularisation and structure hyperparameters (max_depth, min_child_weight, subsample, colsample_bytree, lambda) rank-transfer well across data size. Learning rate and the resulting num_round do not - a rate that converged in 300 rounds on 20% of the data will need more rounds at full scale. Re-tune those two, and with early stopping on a real validation fold, at full size.

Also fix the past mistake explicitly: never deploy a subsample-trained model. The 0.911 is a tuning signal, not a production estimate.

Input mode: File mode copies 200 GB to the instance's EBS volume before training starts. Two traps: the default training volume is far smaller than 200 GB, so the job dies with a disk-full error unless volume_size is set to roughly 400 GB; and you pay the download on every trial. Pipe mode streams from S3 with no disk copy, starts training sooner, and the XGBoost container supports it for CSV and libsvm/protobuf. For HPO, Pipe mode is the clear choice.

Object layout:

• Avoid one giant 200 GB object: a single object is read by one reader and throttles parallelism.
• Avoid tens of thousands of tiny objects: per-object request overhead dominates.
• Target roughly 8-16 shards of 8-16 GB each, so multiple readers and, if you later distribute, multiple workers can read in parallel.
• CSV carries no schema, so a header in every shard confuses the parser; either keep the header only in the first shard or switch to protobuf, which encodes shape and sparsity and is the faster format for this data.

Also set: SageMakerContainerArguments / num_round explicitly, a validation channel for early stopping, and checkpoint_s3_uri if you intend to use managed spot for the tuning trials.

Do the memory arithmetic before deciding. 60M rows x 900 features x 4 bytes ≈ 216 GB if fully dense. The file is 200 GB of CSV, so the matrix must be substantially sparse; XGBoost's DMatrix stays roughly in that range but compressed for sparse columns.

Decision path:

1. Measure actual DMatrix memory on one instance. If it fits on a single memory-optimised instance (ml.r5.24xlarge class has ~768 GB), scale up, not out. Single-node training has no inter-worker communication and is almost always faster and cheaper than data-parallel at this size.
2. If it does not fit, distributed data-parallel splits rows across workers, but every iteration performs an allreduce of gradients/histograms over the network, and the benefit only appears when per-worker data is still large enough to keep CPUs busy. Going from 1 to 8 workers on a shrunken subsample often gets slower.
3. Cheaper wins first: drop or hash low-importance features from the 900 (900 to 400 typically costs a fraction of an AUC point and cuts memory and time proportionally), use sparse libsvm/protobuf input, and raise max_bin carefully rather than blindly.

If you do distribute, use the same instance type for every worker, an even row split, and S3 checkpointing so an interruption does not restart from zero.