Repository: autogluon/autogluon
Pull request: [Tabular] Added use_bag_holdout functionality
*Issue #, if available:*

*Description of changes:*

- Added use_bag_holdout option to allow bagging/stacking with a holdout set. (This is not equivalent to blending)
- Greatly improved handling of rare classes in various situations + bug fixes to rare class duplication logic.

Before merging:

- [x] Add unit test
- [x] Benchmark
 
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.


Problem statement:
Stacking with evaluation set
I'm working with patient data collected for a few years. 
The last year of the data is used as a test set and the year before is kept for validation.

According to the documentation, it is not possible to enable stacking and use a validation set. If stacking is enabled, the validation set is also used for training, I assume the score reported is the out-of-bag score. However, in my case, this badly overfits the data, the score on the test set is low while the out-of-bag score is relatively good, it will be more beneficial in this case to report the score on the validation set and not use it for training, this will enable me to select the best performing model for my needs.

Using hyperparameter tunning instead of stacking and maximizing the score on the validation set gives better performance, which is not recommended in the documentation.

Am I missing something? Perhaps there is a workaround to enable scoring of the validation set?

Thank you all for your great job!

Hints from issue discussion:
Thanks for the feature request! This is a good feature to have. AutoGluon does not currently support specifying a custom validation set for stacking as we use out-of-fold predictions. However, if you have co-variate shift in your data, this can result in heavy overfitting and worse results then if the validation set was used.

Adding this requires updating the stack ensemble API to take in validation data and early-stop/score using it rather than the out-of-fold predictions. This is do-able, but is a fairly complicated effort. I will see about this making it into the 0.1 release as I think it is a very useful feature.

In the meantime, you may be able to get better results by removing the time feature from your data or doing feature engineering on it so that it is simply dayofweek, monthofyear, etc. Tabular models often perform poorly when values are all smaller in training than in validation/test (such as seconds_from_epoch if data is split by time).

Another option is to use your validation data as input to `predictor.leaderboard(val_data)` in order to find the best scoring model on your validation data.

Commit message: minor fix

Patch to review:
diff --git a/core/src/autogluon/core/models/__init__.py b/core/src/autogluon/core/models/__init__.py
index 2f05c4ae40c2..9e74d0e3c994 100644
--- a/core/src/autogluon/core/models/__init__.py
+++ b/core/src/autogluon/core/models/__init__.py
@@ -1,5 +1,5 @@
 from .abstract.abstract_model import AbstractModel
-from .greedy_ensemble.greedy_weighted_ensemble_model import GreedyWeightedEnsembleModel
+from .greedy_ensemble.greedy_weighted_ensemble_model import GreedyWeightedEnsembleModel, SimpleWeightedEnsembleModel
 from .ensemble.bagged_ensemble_model import BaggedEnsembleModel
 from .ensemble.stacker_ensemble_model import StackerEnsembleModel
 from .ensemble.weighted_ensemble_model import WeightedEnsembleModel
diff --git a/core/src/autogluon/core/models/abstract/abstract_model.py b/core/src/autogluon/core/models/abstract/abstract_model.py
index 8a2c7cc677d7..b4c499c0a0c0 100644
--- a/core/src/autogluon/core/models/abstract/abstract_model.py
+++ b/core/src/autogluon/core/models/abstract/abstract_model.py
@@ -144,6 +144,8 @@ def __init__(self,
 
         self.params_trained = dict()
         self._is_initialized = False
+        self._is_fit_metadata_registered = False
+        self._fit_metadata = dict()
 
     def _init_params(self):
         hyperparameters = self._user_params
@@ -182,6 +184,10 @@ def is_fit(self) -> bool:
         """Returns True if the model has been fit."""
         return self.model is not None
 
+    def can_fit(self) -> bool:
+        """Returns True if the model can be fit."""
+        return not self.is_fit()
+
     # TODO: v0.1 update to be aligned with _set_default_auxiliary_params(), add _get_default_params()
     def _set_default_params(self):
         pass
@@ -438,6 +444,21 @@ def _preprocess_fit_resources(self, silent=False, **kwargs):
             logger.log(15, f"\tFitting {self.name} with 'num_gpus': {kwargs['num_gpus']}, 'num_cpus': {kwargs['num_cpus']}")
         return kwargs
 
+    def _register_fit_metadata(self, **kwargs):
+        """
+        Used to track properties of the inputs received during fit, such as if validation data was present.
+        """
+        if not self._is_fit_metadata_registered:
+            self._fit_metadata = self._compute_fit_metadata(**kwargs)
+            self._is_fit_metadata_registered = True
+
+    def _compute_fit_metadata(self, X_val=None, X_unlabeled=None, **kwargs):
+        fit_metadata = dict(
+            val_in_fit=X_val is not None,
+            unlabeled_in_fit=X_unlabeled is not None
+        )
+        return fit_metadata
+
     def fit(self, **kwargs):
         """
         Fit model to predict values in y based on X.
@@ -493,6 +514,7 @@ def fit(self, **kwargs):
         kwargs = self.initialize(**kwargs)  # FIXME: This might have to go before self._preprocess_fit_args, but then time_limit might be incorrect in **kwargs init to initialize
         kwargs = self._preprocess_fit_args(**kwargs)
         if 'time_limit' not in kwargs or kwargs['time_limit'] is None or kwargs['time_limit'] > 0:
+            self._register_fit_metadata(**kwargs)
             self._fit(**kwargs)
         else:
             logger.warning(f'\tWarning: Model has no time left to train, skipping model... (Time Left = {round(kwargs["time_limit"], 1)}s)')
diff --git a/core/src/autogluon/core/models/ensemble/bagged_ensemble_model.py b/core/src/autogluon/core/models/ensemble/bagged_ensemble_model.py
index afcae7a61c58..463e332ef244 100644
--- a/core/src/autogluon/core/models/ensemble/bagged_ensemble_model.py
+++ b/core/src/autogluon/core/models/ensemble/bagged_ensemble_model.py
@@ -84,6 +84,9 @@ def is_stratified(self):
     def is_fit(self):
         return len(self.models) != 0
 
+    def can_fit(self) -> bool:
+        return not self.is_fit() or self._bagged_mode
+
     def is_valid_oof(self):
         return self.is_fit() and (self._child_oof or self._bagged_mode)
 
@@ -112,6 +115,8 @@ def preprocess(self, X, preprocess_nonadaptive=True, model=None, **kwargs):
     def _fit(self,
              X,
              y,
+             X_val=None,
+             y_val=None,
              k_fold=5,
              k_fold_start=0,
              k_fold_end=None,
diff --git a/core/src/autogluon/core/models/greedy_ensemble/ensemble_selection.py b/core/src/autogluon/core/models/greedy_ensemble/ensemble_selection.py
index f3358c34e55a..9bf9f6e5f28b 100644
--- a/core/src/autogluon/core/models/greedy_ensemble/ensemble_selection.py
+++ b/core/src/autogluon/core/models/greedy_ensemble/ensemble_selection.py
@@ -11,7 +11,22 @@
 logger = logging.getLogger(__name__)
 
 
-class EnsembleSelection:
+class AbstractWeightedEnsemble:
+    def predict(self, X):
+        y_pred_proba = self.predict_proba(X)
+        return get_pred_from_proba(y_pred_proba=y_pred_proba, problem_type=self.problem_type)
+
+    def predict_proba(self, X):
+        return self.weight_pred_probas(X, weights=self.weights_)
+
+    @staticmethod
+    def weight_pred_probas(pred_probas, weights):
+        preds_norm = [pred * weight for pred, weight in zip(pred_probas, weights)]
+        preds_ensemble = np.sum(preds_norm, axis=0)
+        return preds_ensemble
+
+
+class EnsembleSelection(AbstractWeightedEnsemble):
     def __init__(
             self,
             ensemble_size: int,
@@ -164,15 +179,13 @@ def _calculate_weights(self):
 
         self.weights_ = weights
 
-    def predict(self, X):
-        y_pred_proba = self.predict_proba(X)
-        return get_pred_from_proba(y_pred_proba=y_pred_proba, problem_type=self.problem_type)
 
-    def predict_proba(self, X):
-        return self.weight_pred_probas(X, weights=self.weights_)
+class SimpleWeightedEnsemble(AbstractWeightedEnsemble):
+    """Predefined user-weights ensemble"""
+    def __init__(self, weights, problem_type, **kwargs):
+        self.weights_ = weights
+        self.problem_type = problem_type
 
-    @staticmethod
-    def weight_pred_probas(pred_probas, weights):
-        preds_norm = [pred * weight for pred, weight in zip(pred_probas, weights)]
-        preds_ensemble = np.sum(preds_norm, axis=0)
-        return preds_ensemble
+    @property
+    def ensemble_size(self):
+        return len(self.weights_)
diff --git a/core/src/autogluon/core/models/greedy_ensemble/greedy_weighted_ensemble_model.py b/core/src/autogluon/core/models/greedy_ensemble/greedy_weighted_ensemble_model.py
index b55252a551f0..19c007b0a288 100644
--- a/core/src/autogluon/core/models/greedy_ensemble/greedy_weighted_ensemble_model.py
+++ b/core/src/autogluon/core/models/greedy_ensemble/greedy_weighted_ensemble_model.py
@@ -1,6 +1,6 @@
 import logging
 
-from .ensemble_selection import EnsembleSelection
+from .ensemble_selection import EnsembleSelection, SimpleWeightedEnsemble
 from ..abstract.abstract_model import AbstractModel
 from ..abstract.model_trial import skip_hpo
 from ...constants import MULTICLASS, SOFTCLASS, QUANTILE
@@ -56,10 +56,12 @@ def _fit(self,
              time_limit=None,
              sample_weight=None,
              **kwargs):
-        X = self.preprocess(X)
-        self.model = self.model_base(ensemble_size=self.params['ensemble_size'], problem_type=self.problem_type, quantile_levels=self.quantile_levels, metric=self.stopping_metric)
-        self.model = self.model.fit(X, y, time_limit=time_limit, sample_weight=sample_weight)
-        self.base_model_names, self.model.weights_ = self.remove_zero_weight_models(self.base_model_names, self.model.weights_)
+        params = self._get_model_params()
+        if self.model is None:
+            X = self.preprocess(X)
+            self.model = self.model_base(problem_type=self.problem_type, quantile_levels=self.quantile_levels, metric=self.stopping_metric, **params)
+            self.model = self.model.fit(X, y, time_limit=time_limit, sample_weight=sample_weight)
+            self.base_model_names, self.model.weights_ = self.remove_zero_weight_models(self.base_model_names, self.model.weights_)
         self.features = self._set_stack_columns(base_model_names=self.base_model_names)
         self.params_trained['ensemble_size'] = self.model.ensemble_size
         self.weights_ = self.model.weights_
@@ -131,3 +133,21 @@ def _get_default_stopping_metric(self):
 
     def _hyperparameter_tune(self, **kwargs):
         return skip_hpo(self, **kwargs)
+
+
+class SimpleWeightedEnsembleModel(GreedyWeightedEnsembleModel):
+    def __init__(self, model_base=SimpleWeightedEnsemble, **kwargs):
+        super().__init__(model_base=model_base, **kwargs)
+
+    def _fit(self, **kwargs):
+        params = self._get_model_params()
+        if 'weights' not in params:
+            raise ValueError('Missing required parameter "weights" to fit SimpleWeightedEnsembleModel.')
+        if len(params['weights']) != len(self.base_model_names):
+            raise AssertionError(f'Length of weights does not equal length of self.base_model_names ({len(params["weights"])}, {len(self.base_model_names)})')
+
+        if self.model is None:
+            self.model = self.model_base(problem_type=self.problem_type, **params)
+            self.base_model_names, self.model.weights_ = self.remove_zero_weight_models(self.base_model_names, self.model.weights_)
+        self.features = self._set_stack_columns(base_model_names=self.base_model_names)
+        self.weights_ = self.model.weights_
diff --git a/core/src/autogluon/core/utils/utils.py b/core/src/autogluon/core/utils/utils.py
index 8c4bd729ab70..bf41c34e07ec 100644
--- a/core/src/autogluon/core/utils/utils.py
+++ b/core/src/autogluon/core/utils/utils.py
@@ -5,6 +5,7 @@
 import math
 import pickle
 import time
+import random
 import sys
 from typing import Callable
 from datetime import datetime
@@ -18,6 +19,7 @@
 from sklearn.model_selection import KFold, StratifiedKFold, RepeatedKFold, RepeatedStratifiedKFold
 from sklearn.model_selection import train_test_split
 
+from .miscs import warning_filter
 from ..constants import BINARY, REGRESSION, MULTICLASS, SOFTCLASS, QUANTILE
 from ..metrics import accuracy, root_mean_squared_error, pinball_loss, Scorer
 from ..features.infer_types import get_type_map_raw
@@ -47,7 +49,24 @@ def generate_kfold(X, y=None, n_splits=5, random_state=0, stratified=False, n_re
             kf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
 
         kf.get_n_splits(X, y)
-        return [[train_index, test_index] for train_index, test_index in kf.split(X, y)]
+        # FIXME: There is a bug in sklearn that causes an incorrect ValueError if performing stratification and all classes have fewer than n_splits samples.
+        #  This is hacked by adding a dummy class with n_splits samples, performing the kfold split, then removing the dummy samples from all resulting indices.
+        #  This is very inefficient and complicated and ideally should be fixed in sklearn.
+        with warning_filter():
+            try:
+                out = [[train_index, test_index] for train_index, test_index in kf.split(X, y)]
+            except:
+                y_dummy = pd.concat([y, pd.Series([-1] * n_splits)], ignore_index=True)
+                X_dummy = pd.concat([X, X.head(n_splits)], ignore_index=True)
+                invalid_index = set(list(y_dummy.tail(n_splits).index))
+                out = [[train_index, test_index] for train_index, test_index in kf.split(X_dummy, y_dummy)]
+                len_out = len(out)
+                for i in range(len_out):
+                    train_index = out[i][0]
+                    test_index = out[i][1]
+                    out[i][0] = [index for index in train_index if index not in invalid_index]
+                    out[i][1] = [index for index in test_index if index not in invalid_index]
+        return out
     else:
         if n_repeats > 1:
             kf = RepeatedKFold(n_splits=n_splits, n_repeats=n_repeats, random_state=random_state)
@@ -240,13 +259,21 @@ def augment_rare_classes(X, label, threshold):
             new_df = new_df.append(clss_df.copy())
         aug_df = aug_df.append(new_df.copy())
 
+    # Ensure new samples generated via augmentation have unique indices
+    aug_df = aug_df.reset_index(drop=True)
+    aug_df_len = len(aug_df)
+    X_index_aug_start = max(list(X.index)) + 1
+    aug_index = [X_index_aug_start + i for i in range(aug_df_len)]
+    aug_df.index = aug_index
+
+    logger.log(20, f"Duplicated {len(aug_df)} samples from {len(class_counts_invalid)} rare classes in training set because eval_metric requires all classes have at least {threshold} samples.")
+
     X = X.append(aug_df)
     class_counts = X[label].value_counts()
     class_counts_invalid = class_counts[class_counts < threshold]
     class_counts_invalid = class_counts_invalid[~class_counts_invalid.index.isin(set(missing_classes))]
     if len(class_counts_invalid) > 0:
         raise RuntimeError("augment_rare_classes failed to produce enough data from rare classes")
-    logger.log(15, "Replicated some data from rare classes in training set because eval_metric requires all classes")
     return X
 
 
@@ -273,26 +300,82 @@ def get_pred_from_proba(y_pred_proba, problem_type=BINARY):
     return y_pred
 
 
-def generate_train_test_split(X: DataFrame, y: Series, problem_type: str, test_size: float = 0.1, random_state=0) -> (DataFrame, DataFrame, Series, Series):
+def generate_train_test_split(X: DataFrame,
+                              y: Series,
+                              problem_type: str,
+                              test_size: float = 0.1,
+                              random_state=0,
+                              min_cls_count_train=1) -> (DataFrame, DataFrame, Series, Series):
     if (test_size <= 0.0) or (test_size >= 1.0):
         raise ValueError("fraction of data to hold-out must be specified between 0 and 1")
 
-    if problem_type in [REGRESSION, QUANTILE, SOFTCLASS]:
-        stratify = None
-    else:
+    X_split = X
+    y_split = y
+    if problem_type in [BINARY, MULTICLASS]:
         stratify = y
+    else:
+        stratify = None
 
-    # TODO: Enable stratified split when y class would result in 0 samples in test.
-    #  One approach: extract low frequency classes from X/y, add back (1-test_size)% to X_train, y_train, rest to X_test
-    #  Essentially stratify the high frequency classes, random the low frequency (While ensuring at least 1 example stays for each low frequency in train!)
-    #  Alternatively, don't test low frequency at all, trust it to work in train set. Risky, but highest quality for predictions.
-    X_train, X_test, y_train, y_test = train_test_split(X, y.values, test_size=test_size, shuffle=True, random_state=random_state, stratify=stratify)
+    # This code block is necessary to avoid crashing when performing a stratified split and only 1 sample exists for a class.
+    # This code will ensure that the sample will be part of the train split, meaning the test split will have 0 samples of the rare class.
+    rare_indices = None
+    if stratify is not None:
+        cls_counts = y.value_counts()
+        cls_counts_invalid = cls_counts[cls_counts < min_cls_count_train]
+
+        if len(cls_counts_invalid) > 0:
+            logger.error(f'ERROR: Classes have too few samples to split the data! At least {min_cls_count_train} are required.')
+            logger.error(cls_counts_invalid)
+            raise AssertionError('Not enough data to split data into train and val without dropping classes!')
+        elif min_cls_count_train < 2:
+            cls_counts_rare = cls_counts[cls_counts < 2]
+            if len(cls_counts_rare) > 0:
+                cls_counts_rare_val = set(cls_counts_rare.index)
+                y_rare = y[y.isin(cls_counts_rare_val)]
+                rare_indices = list(y_rare.index)
+                X_split = X.drop(index=rare_indices)
+                y_split = y.drop(index=rare_indices)
+                stratify = y_split
+
+    X_train, X_test, y_train, y_test = train_test_split(X_split, y_split.values, test_size=test_size, shuffle=True, random_state=random_state, stratify=stratify)
     if problem_type != SOFTCLASS:
         y_train = pd.Series(y_train, index=X_train.index)
         y_test = pd.Series(y_test, index=X_test.index)
     else:
         y_train = pd.DataFrame(y_train, index=X_train.index)
         y_test = pd.DataFrame(y_test, index=X_test.index)
+
+    if rare_indices:
+        X_train = X_train.append(X.loc[rare_indices])
+        y_train = y_train.append(y.loc[rare_indices])
+
+    if problem_type in [BINARY, MULTICLASS]:
+        class_counts_dict_orig = y.value_counts().to_dict()
+        class_counts_dict = y_train.value_counts().to_dict()
+        class_counts_dict_test = y_test.value_counts().to_dict()
+
+        indices_to_move = []
+        random_state = random.getstate()
+        random.seed(random_state)
+        for cls in class_counts_dict_orig.keys():
+            count = class_counts_dict.get(cls, 0)
+            if count >= min_cls_count_train:
+                continue
+            count_test = class_counts_dict_test.get(cls, 0)
+            if count + count_test < min_cls_count_train:
+                raise AssertionError('Not enough data to split data into train and val without dropping classes!')
+            count_to_move = min_cls_count_train - count
+            indices_of_cls_test = list(y_test[y_test == cls].index)
+            indices_to_move_cls = random.sample(indices_of_cls_test, count_to_move)
+            indices_to_move += indices_to_move_cls
+        random.setstate(random_state)
+        if indices_to_move:
+            y_test_moved = y_test.loc[indices_to_move].copy()
+            X_test_moved = X_test.loc[indices_to_move].copy()
+            y_train = pd.concat([y_train, y_test_moved])
+            X_train = pd.concat([X_train, X_test_moved])
+            y_test = y_test.drop(index=indices_to_move)
+            X_test = X_test.drop(index=indices_to_move)
     return X_train, X_test, y_train, y_test
 
 
diff --git a/tabular/src/autogluon/tabular/learner/default_learner.py b/tabular/src/autogluon/tabular/learner/default_learner.py
index 3c90a0f69452..776af029326a 100644
--- a/tabular/src/autogluon/tabular/learner/default_learner.py
+++ b/tabular/src/autogluon/tabular/learner/default_learner.py
@@ -127,16 +127,20 @@ def general_data_processing(self, X: DataFrame, X_val: DataFrame, X_unlabeled: D
 
         if X_val is not None and self.label in X_val.columns:
             holdout_frac = 1
-        self.threshold, holdout_frac, num_bag_folds = self.adjust_threshold_if_necessary(X[self.label], threshold=self.threshold, holdout_frac=holdout_frac, num_bag_folds=num_bag_folds)
 
         if (self.eval_metric is not None) and (self.eval_metric.name in ['log_loss', 'pac_score']) and (self.problem_type == MULTICLASS):
-            X = augment_rare_classes(X, self.label, self.threshold)
+            if num_bag_folds > 0:
+                self.threshold = 2
+                X = augment_rare_classes(X, self.label, threshold=2)
+            else:
+                self.threshold = 1
+
+        self.threshold, holdout_frac, num_bag_folds = self.adjust_threshold_if_necessary(X[self.label], threshold=self.threshold, holdout_frac=holdout_frac, num_bag_folds=num_bag_folds)
 
         # Gets labels prior to removal of infrequent classes
         y_uncleaned = X[self.label].copy()
 
         self.cleaner = Cleaner.construct(problem_type=self.problem_type, label=self.label, threshold=self.threshold)
-        # TODO: What if all classes in X are low frequency in multiclass? Currently we would crash. Not certain how many problems actually have this property
         X = self.cleaner.fit_transform(X)  # TODO: Consider merging cleaner into label_cleaner
         X, y = self.extract_label(X)
         self.label_cleaner = LabelCleaner.construct(problem_type=self.problem_type, y=y, y_uncleaned=y_uncleaned, positive_class=self._positive_class)
@@ -254,20 +258,21 @@ def adjust_threshold_if_necessary(self, y, threshold, holdout_frac, num_bag_fold
 
     def _adjust_threshold_if_necessary(self, y, threshold, holdout_frac, num_bag_folds):
         new_threshold = threshold
-        if self.problem_type in [REGRESSION, QUANTILE]:
-            num_rows = len(y)
-            holdout_frac = max(holdout_frac, 1 / num_rows + 0.001)
-            num_bag_folds = min(num_bag_folds, num_rows)
-            return new_threshold, holdout_frac, num_bag_folds
+        num_rows = len(y)
+        holdout_frac = max(holdout_frac, 1 / num_rows + 0.001)
+        num_bag_folds = min(num_bag_folds, num_rows)
 
         if num_bag_folds < 2:
-            minimum_safe_threshold = math.ceil(1 / holdout_frac)
+            minimum_safe_threshold = 1
         else:
-            minimum_safe_threshold = num_bag_folds
+            minimum_safe_threshold = 2
 
         if minimum_safe_threshold > new_threshold:
             new_threshold = minimum_safe_threshold
 
+        if self.problem_type in [REGRESSION, QUANTILE]:
+            return new_threshold, holdout_frac, num_bag_folds
+
         class_counts = y.value_counts()
         total_rows = class_counts.sum()
         minimum_percent_to_keep = 0.975
@@ -292,17 +297,6 @@ def _adjust_threshold_if_necessary(self, y, threshold, holdout_frac, num_bag_fol
             if (num_rows_valid >= minimum_rows_to_keep) and (num_classes_valid >= minimum_class_to_keep):
                 break
 
-        if new_threshold == 1:
-            new_threshold = 2  # threshold=1 is invalid, can't perform any train/val split in this case.
-        self.threshold = new_threshold
-
-        if new_threshold < minimum_safe_threshold:
-            if num_bag_folds >= 2:
-                if num_bag_folds > new_threshold:
-                    num_bag_folds = new_threshold
-            elif math.ceil(1 / holdout_frac) > new_threshold:
-                holdout_frac = 1 / new_threshold + 0.001
-
         return new_threshold, holdout_frac, num_bag_folds
 
     def get_info(self, include_model_info=False, **kwargs):
diff --git a/tabular/src/autogluon/tabular/predictor/predictor.py b/tabular/src/autogluon/tabular/predictor/predictor.py
index 8b1fbaefd928..1a5f7341afd4 100644
--- a/tabular/src/autogluon/tabular/predictor/predictor.py
+++ b/tabular/src/autogluon/tabular/predictor/predictor.py
@@ -471,10 +471,16 @@ def fit(self,
                 Disabled by default (0), but we recommend values between 1-3 to maximize predictive performance.
                 To prevent overfitting, `num_bag_folds >= 2` must also be set or else a ValueError will be raised.
             holdout_frac : float, default = None
-                Fraction of train_data to holdout as tuning data for optimizing hyperparameters (ignored unless `tuning_data = None`, ignored if `num_bag_folds != 0`).
+                Fraction of train_data to holdout as tuning data for optimizing hyperparameters (ignored unless `tuning_data = None`, ignored if `num_bag_folds != 0` unless `use_bag_holdout == True`).
                 Default value (if None) is selected based on the number of rows in the training data. Default values range from 0.2 at 2,500 rows to 0.01 at 250,000 rows.
                 Default value is doubled if `hyperparameter_tune_kwargs` is set, up to a maximum of 0.2.
-                Disabled if `num_bag_folds >= 2`.
+                Disabled if `num_bag_folds >= 2` unless `use_bag_holdout == True`.
+            use_bag_holdout : bool, default = False
+                If True, a `holdout_frac` portion of the data is held-out from model bagging.
+                This held-out data is only used to score models and determine weighted ensemble weights.
+                Enable this if there is a large gap between score_val and score_test in stack models.
+                Note: If `tuning_data` was specified, `tuning_data` is used as the holdout data.
+                Disabled if not bagging.
             hyperparameter_tune_kwargs : str or dict, default = None
                 Hyperparameter tuning strategy and kwargs (for example, how many HPO trials to run).
                 If None, then hyperparameter tuning will not be performed.
@@ -644,6 +650,7 @@ def fit(self,
         ag_args_fit = kwargs['ag_args_fit']
         ag_args_ensemble = kwargs['ag_args_ensemble']
         excluded_model_types = kwargs['excluded_model_types']
+        use_bag_holdout = kwargs['use_bag_holdout']
 
         if ag_args is None:
             ag_args = {}
@@ -705,7 +712,7 @@ def fit(self,
         }
         self._learner.fit(X=train_data, X_val=tuning_data, X_unlabeled=unlabeled_data,
                           holdout_frac=holdout_frac, num_bag_folds=num_bag_folds, num_bag_sets=num_bag_sets, num_stack_levels=num_stack_levels,
-                          hyperparameters=hyperparameters, core_kwargs=core_kwargs, time_limit=time_limit, verbosity=verbosity)
+                          hyperparameters=hyperparameters, core_kwargs=core_kwargs, time_limit=time_limit, verbosity=verbosity, use_bag_holdout=use_bag_holdout)
         self._set_post_fit_vars()
 
         self._post_fit(
@@ -1708,10 +1715,16 @@ def get_oof_pred_proba(self, model: str = None, transformed=False, as_multiclass
         -------
         :class:`pd.Series` or :class:`pd.DataFrame` object of the out-of-fold training prediction probabilities of the model.
         """
-        if not self._trainer.bagged_mode:
-            raise AssertionError('Predictor must be in bagged mode to get out-of-fold predictions.')
         if model is None:
             model = self.get_model_best()
+        if not self._trainer.bagged_mode:
+            raise AssertionError('Predictor must be in bagged mode to get out-of-fold predictions.')
+        if model in self._trainer._model_full_dict_val_score:
+            # FIXME: This is a hack, add refit tag in a nicer way than via the _model_full_dict_val_score
+            # TODO: bagged-with-holdout refit to bagged-no-holdout should still be able to return out-of-fold predictions
+            raise AssertionError('_FULL models do not have out-of-fold predictions.')
+        if self._trainer.get_model_attribute_full(model=model, attribute='val_in_fit', func=max):
+            raise AssertionError(f'Model {model} does not have out-of-fold predictions because it used a validation set during training.')
         y_pred_proba_oof_transformed = self.transform_features(base_models=[model], return_original_features=False)
         if not internal_oof:
             is_duplicate_index = y_pred_proba_oof_transformed.index.duplicated(keep='first')
@@ -2204,6 +2217,7 @@ def _validate_fit_kwargs(self, kwargs):
             holdout_frac=None,  # TODO: Potentially error if num_bag_folds is also specified
             num_bag_folds=None,  # TODO: Potentially move to fit_extra, raise exception if value too large / invalid in fit_extra.
             auto_stack=False,
+            use_bag_holdout=False,
 
             # other
             feature_generator='auto',
diff --git a/tabular/src/autogluon/tabular/trainer/abstract_trainer.py b/tabular/src/autogluon/tabular/trainer/abstract_trainer.py
index f3d22134a9cd..2b0ca80a813f 100644
--- a/tabular/src/autogluon/tabular/trainer/abstract_trainer.py
+++ b/tabular/src/autogluon/tabular/trainer/abstract_trainer.py
@@ -8,7 +8,7 @@
 import psutil
 from collections import defaultdict
 
-from autogluon.core.constants import AG_ARGS_FIT, BINARY, MULTICLASS, REGRESSION, QUANTILE, REFIT_FULL_NAME, REFIT_FULL_SUFFIX
+from autogluon.core.constants import AG_ARGS, AG_ARGS_FIT, BINARY, MULTICLASS, REGRESSION, QUANTILE, REFIT_FULL_NAME, REFIT_FULL_SUFFIX
 from autogluon.core.models import AbstractModel, BaggedEnsembleModel, StackerEnsembleModel, WeightedEnsembleModel
 from autogluon.core.features.feature_metadata import FeatureMetadata
 from autogluon.core.scheduler.scheduler_factory import scheduler_factory
@@ -85,6 +85,10 @@ def __init__(self, path: str, problem_type: str, eval_metric=None,
         self._num_cols_train = None
 
         self.is_data_saved = False
+        self._X_saved = False
+        self._y_saved = False
+        self._X_val_saved = False
+        self._y_val_saved = False
 
         self._regress_preds_asprobas = False  # whether to treat regression predictions as class-probabilities (during distillation)
 
@@ -106,48 +110,56 @@ def path_data(self) -> str:
         return self.path_utils + 'data' + os.path.sep
 
     def load_X(self):
-        path = self.path_data + 'X.pkl'
-        return load_pkl.load(path=path)
+        if self._X_saved:
+            path = self.path_data + 'X.pkl'
+            return load_pkl.load(path=path)
+        return None
 
     def load_X_val(self):
-        path = self.path_data + 'X_val.pkl'
-        return load_pkl.load(path=path)
+        if self._X_val_saved:
+            path = self.path_data + 'X_val.pkl'
+            return load_pkl.load(path=path)
+        return None
 
     def load_y(self):
-        path = self.path_data + 'y.pkl'
-        return load_pkl.load(path=path)
+        if self._y_saved:
+            path = self.path_data + 'y.pkl'
+            return load_pkl.load(path=path)
+        return None
 
     def load_y_val(self):
-        path = self.path_data + 'y_val.pkl'
-        return load_pkl.load(path=path)
+        if self._y_val_saved:
+            path = self.path_data + 'y_val.pkl'
+            return load_pkl.load(path=path)
+        return None
 
     def load_data(self):
         X = self.load_X()
         y = self.load_y()
-        if not self.bagged_mode:
-            X_val = self.load_X_val()
-            y_val = self.load_y_val()
-        else:
-            X_val = None
-            y_val = None
+        X_val = self.load_X_val()
+        y_val = self.load_y_val()
 
         return X, y, X_val, y_val
 
     def save_X(self, X, verbose=True):
         path = self.path_data + 'X.pkl'
         save_pkl.save(path=path, object=X, verbose=verbose)
+        self._X_saved = True
 
     def save_X_val(self, X, verbose=True):
         path = self.path_data + 'X_val.pkl'
         save_pkl.save(path=path, object=X, verbose=verbose)
+        self._X_val_saved = True
 
     def save_y(self, y, verbose=True):
         path = self.path_data + 'y.pkl'
         save_pkl.save(path=path, object=y, verbose=verbose)
+        self._y_saved = True
 
     def save_y_val(self, y, verbose=True):
         path = self.path_data + 'y_val.pkl'
         save_pkl.save(path=path, object=y, verbose=verbose)
+        self._y_val_saved = True
 
     def get_model_names(self, stack_name: Union[List[str], str] = None, level: Union[List[int], int] = None, can_infer: bool = None, models: List[str] = None) -> List[str]:
         if models is None:
@@ -301,7 +313,7 @@ def stack_new_level(self, X, y, models: Union[List[AbstractModel], dict], X_val=
         core_models = self.stack_new_level_core(X=X, y=y, X_val=X_val, y_val=y_val, X_unlabeled=X_unlabeled, models=models,
                                                 level=level, base_model_names=base_model_names, feature_prune=feature_prune, **core_kwargs)
 
-        if self.bagged_mode:
+        if X_val is None:
             aux_models = self.stack_new_level_aux(X=X, y=y, base_model_names=core_models, level=level+1, **aux_kwargs)
         else:
             aux_models = self.stack_new_level_aux(X=X_val, y=y_val, fit=False, base_model_names=core_models, level=level+1, **aux_kwargs)
@@ -310,7 +322,7 @@ def stack_new_level(self, X, y, models: Union[List[AbstractModel], dict], X_val=
     def stack_new_level_core(self, X, y, models: Union[List[AbstractModel], dict], X_val=None, y_val=None, X_unlabeled=None,
                              level=1, base_model_names: List[str] = None, stack_name='core',
                              ag_args=None, ag_args_fit=None, ag_args_ensemble=None, excluded_model_types=None, ensemble_type=StackerEnsembleModel,
-                             name_suffix: str = None, get_models_func=None, **kwargs) -> List[str]:
+                             name_suffix: str = None, get_models_func=None, refit_full=False, **kwargs) -> List[str]:
         """
         Trains all models using the data provided.
         If level > 1, then the models will use base model predictions as additional features.
@@ -367,11 +379,19 @@ def stack_new_level_core(self, X, y, models: Union[List[AbstractModel], dict], X
         X_init = self.get_inputs_to_stacker(X, base_models=base_model_names, fit=True)
         if X_val is not None:
             X_val = self.get_inputs_to_stacker(X_val, base_models=base_model_names, fit=False)
+        if refit_full and X_val is not None:
+            X_init = pd.concat([X_init, X_val])
+            y = pd.concat([y, y_val])
+            X_val = None
+            y_val = None
         if X_unlabeled is not None:
             X_unlabeled = self.get_inputs_to_stacker(X_unlabeled, base_models=base_model_names, fit=False)
 
+        fit_kwargs = dict(num_classes=self.num_classes)
+
         # FIXME: TODO: v0.1 X_unlabeled isn't cached so it won't be available during refit_full or fit_extra.
-        return self._train_multi(X=X_init, y=y, X_val=X_val, y_val=y_val, X_unlabeled=X_unlabeled, models=models, level=level, stack_name=stack_name, **kwargs)
+        return self._train_multi(X=X_init, y=y, X_val=X_val, y_val=y_val, X_unlabeled=X_unlabeled,
+                                 models=models, level=level, stack_name=stack_name, fit_kwargs=fit_kwargs, **kwargs)
 
     # TODO: Consider making level be auto-determined based off of max(base_model_levels)+1
     # TODO: Remove name_suffix, hacked in
@@ -577,19 +597,12 @@ def _get_inputs_to_stacker_legacy(self, X, level_start, level_end, model_levels,
     def refit_single_full(self, X=None, y=None, X_val=None, y_val=None, X_unlabeled=None, models=None) -> List[str]:
         if X is None:
             X = self.load_X()
-            if X_val is None and not self.bagged_mode:
-                X_val = self.load_X_val()
+        if X_val is None:
+            X_val = self.load_X_val()
         if y is None:
             y = self.load_y()
-            if y_val is None and not self.bagged_mode:
-                y_val = self.load_y_val()
-
-        if X_val is not None and y_val is not None:
-            X_full = pd.concat([X, X_val])
-            y_full = pd.concat([y, y_val])
-        else:
-            X_full = X
-            y_full = y
+        if y_val is None:
+            y_val = self.load_y_val()
 
         if models is None:
             models = self.get_model_names()
@@ -623,7 +636,7 @@ def refit_single_full(self, X=None, y=None, X_val=None, y_val=None, X_unlabeled=
                 if issubclass(stacker_type, WeightedEnsembleModel):
                     # TODO: Technically we don't need to re-train the weighted ensemble, we could just copy the original and re-use the weights.
                     w = None
-                    if self.bagged_mode:
+                    if X_val is None:
                         if self.weight_evaluation:
                             X, w = extract_column(X, self.sample_weight)
                         X_stack_preds = self.get_inputs_to_stacker(X, base_models=base_model_names, fit=True, use_orig_features=False)
@@ -636,14 +649,19 @@ def refit_single_full(self, X=None, y=None, X_val=None, y_val=None, X_unlabeled=
                     if w is not None:
                         X_stack_preds[self.sample_weight] = w.values/w.mean()
 
-                    # TODO: Remove child_hyperparameters, make this cleaner
-                    #  This fixes the following: Use the original weighted ensemble's iterations: Currently Dionis spends over 1hr training the refit weighted ensemble because it isn't time limited and goes to 100 iterations.
-                    model_child = model_full._get_model_base()
-                    child_hyperparameters = copy.deepcopy(model_child.params)
-                    child_hyperparameters[AG_ARGS_FIT] = copy.deepcopy(model_child.params_aux)
+                    orig_weights = model._get_model_weights()
+                    base_model_names = list(orig_weights.keys())
+                    weights = list(orig_weights.values())
+
+                    child_hyperparameters = {
+                        AG_ARGS: {'model_type': 'SIMPLE_ENS_WEIGHTED'},
+                        'weights': weights,
+                    }
+
                     # TODO: stack_name=REFIT_FULL_NAME_AUX?
                     models_trained = self.generate_weighted_ensemble(X=X_stack_preds, y=y_input, level=level, stack_name=REFIT_FULL_NAME, k_fold=0, n_repeats=1,
-                                                                     base_model_names=base_model_names, name_suffix=REFIT_FULL_SUFFIX, save_bag_folds=True, check_if_best=False, child_hyperparameters=child_hyperparameters)
+                                                                     base_model_names=base_model_names, name_suffix=REFIT_FULL_SUFFIX, save_bag_folds=True,
+                                                                     check_if_best=False, child_hyperparameters=child_hyperparameters)
                     # TODO: Do the below more elegantly, ideally as a parameter to the trainer train function to disable recording scores/pred time.
                     for model_weighted_ensemble in models_trained:
                         model_loaded = self.load_model(model_weighted_ensemble)
@@ -652,8 +670,8 @@ def refit_single_full(self, X=None, y=None, X_val=None, y_val=None, X_unlabeled=
                         self.set_model_attribute(model=model_weighted_ensemble, attribute='val_score', val=None)
                         self.save_model(model_loaded)
                 else:
-                    models_trained = self.stack_new_level_core(X=X_full, y=y_full, X_unlabeled=X_unlabeled, models=[model_full], base_model_names=base_model_names, level=level, stack_name=REFIT_FULL_NAME,
-                                                               hyperparameter_tune_kwargs=None, feature_prune=False, k_fold=0, n_repeats=1, ensemble_type=stacker_type)
+                    models_trained = self.stack_new_level_core(X=X, y=y, X_val=X_val, y_val=y_val, X_unlabeled=X_unlabeled, models=[model_full], base_model_names=base_model_names, level=level, stack_name=REFIT_FULL_NAME,
+                                                               hyperparameter_tune_kwargs=None, feature_prune=False, k_fold=0, n_repeats=1, ensemble_type=stacker_type, refit_full=True)
                 if len(models_trained) == 1:
                     model_full_dict[model_name] = models_trained[0]
                 for model_trained in models_trained:
@@ -834,7 +852,9 @@ def unpersist_models(self, model_names='all') -> list:
             logger.log(30, f'No valid persisted models were specified to be unpersisted, so no change in model persistence was performed.')
         return unpersisted_models
 
-    def generate_weighted_ensemble(self, X, y, level, base_model_names, k_fold=0, n_repeats=1, stack_name=None, hyperparameters=None, time_limit=None, name_suffix: str = None, save_bag_folds=None, check_if_best=True, child_hyperparameters=None, get_models_func=None) -> List[str]:
+    def generate_weighted_ensemble(self, X, y, level, base_model_names, k_fold=0, n_repeats=1, stack_name=None, hyperparameters=None,
+                                   time_limit=None, name_suffix: str = None, save_bag_folds=None, check_if_best=True, child_hyperparameters=None,
+                                   get_models_func=None) -> List[str]:
         if get_models_func is None:
             get_models_func = self.construct_model_templates
         if len(base_model_names) == 0:
@@ -876,7 +896,22 @@ def generate_weighted_ensemble(self, X, y, level, base_model_names, k_fold=0, n_
         w = None
         if self.weight_evaluation:
             X, w = extract_column(X, self.sample_weight)
-        models = self._train_multi(X=X, y=y, X_val=None, y_val=None, models=[weighted_ensemble_model], k_fold=k_fold, n_repeats=n_repeats, hyperparameter_tune_kwargs=None, feature_prune=False, stack_name=stack_name, level=level, time_limit=time_limit, ens_sample_weight=w)
+        models = self._train_multi(
+            X=X,
+            y=y,
+            X_val=None,
+            y_val=None,
+            models=[weighted_ensemble_model],
+            k_fold=k_fold,
+            n_repeats=n_repeats,
+            hyperparameter_tune_kwargs=None,
+            feature_prune=False,
+            stack_name=stack_name,
+            level=level,
+            time_limit=time_limit,
+            ens_sample_weight=w,
+            fit_kwargs=dict(num_classes=self.num_classes),  # FIXME: Is this the right way to do this?
+        )
         for weighted_ensemble_model_name in models:
             if check_if_best and weighted_ensemble_model_name in self.get_model_names():
                 if self.model_best is None:
@@ -894,10 +929,7 @@ def _train_single(self, X, y, model: AbstractModel, X_val=None, y_val=None, **mo
         Trains model but does not add the trained model to this Trainer.
         Returns trained model object.
         """
-        if isinstance(model, BaggedEnsembleModel):
-            model.fit(X=X, y=y, **model_fit_kwargs)
-        else:
-            model.fit(X=X, y=y, X_val=X_val, y_val=y_val, **model_fit_kwargs)
+        model.fit(X=X, y=y, X_val=X_val, y_val=y_val, **model_fit_kwargs)
         return model
 
     def _train_and_save(self, X, y, model: AbstractModel, X_val=None, y_val=None, stack_name='core', level=1, **model_fit_kwargs) -> List[str]:
@@ -934,7 +966,9 @@ def _train_and_save(self, X, y, model: AbstractModel, X_val=None, y_val=None, st
                 w = None
                 w_val = None
             if isinstance(model, BaggedEnsembleModel):
-                if model.is_valid_oof() or isinstance(model, WeightedEnsembleModel):
+                if X_val is not None and y_val is not None:
+                    score = model.score(X=X_val, y=y_val, sample_weight=w_val)
+                elif model.is_valid_oof() or isinstance(model, WeightedEnsembleModel):
                     score = model.score_with_oof(y=y, sample_weight=w)
                 else:
                     score = None
@@ -1031,9 +1065,11 @@ def _add_model(self, model: AbstractModel, stack_name: str = 'core', level: int
             type=type(model),  # Outer type, can be BaggedEnsemble, StackEnsemble (Type that is able to load the model)
             type_inner=type_inner,  # Inner type, if Ensemble then it is the type of the inner model (May not be able to load with this type)
             can_infer=model.can_infer(),
+            can_fit=model.can_fit(),
             is_valid=model.is_valid(),
             stack_name=stack_name,
             level=level,
+            **model._fit_metadata,
         )
         if isinstance(model, StackerEnsembleModel):
             prior_models = self.get_model_names()
@@ -1051,7 +1087,7 @@ def _add_model(self, model: AbstractModel, stack_name: str = 'core', level: int
 
     # TODO: Split this to avoid confusion, HPO should go elsewhere?
     def _train_single_full(self, X, y, model: AbstractModel, X_unlabeled=None, X_val=None, y_val=None, feature_prune=False, hyperparameter_tune_kwargs=None,
-                           stack_name='core', k_fold=None, k_fold_start=0, k_fold_end=None, n_repeats=None, n_repeat_start=0, level=1, time_limit=None, **kwargs) -> List[str]:
+                           stack_name='core', k_fold=None, k_fold_start=0, k_fold_end=None, n_repeats=None, n_repeat_start=0, level=1, time_limit=None, fit_kwargs=None, **kwargs) -> List[str]:
         """
         Trains a model, with the potential to train multiple versions of this model with hyperparameter tuning and feature pruning.
         Returns a list of successfully trained and saved model names.
@@ -1061,10 +1097,13 @@ def _train_single_full(self, X, y, model: AbstractModel, X_unlabeled=None, X_val
             k_fold = self.k_fold
         if n_repeats is None:
             n_repeats = self.n_repeats
+        if fit_kwargs is None:
+            fit_kwargs = dict()
         model_fit_kwargs = dict(
             time_limit=time_limit,
             verbosity=self.verbosity,
         )
+        model_fit_kwargs.update(fit_kwargs)
         if self.sample_weight is not None:
             X, w_train = extract_column(X, self.sample_weight)
             if w_train is not None:  # may be None for ensemble
@@ -1074,7 +1113,7 @@ def _train_single_full(self, X, y, model: AbstractModel, X_unlabeled=None, X_val
                 X_val, w_val = extract_column(X_val, self.sample_weight)
                 if self.weight_evaluation and w_val is not None:  # ignore validation sample weights unless weight_evaluation specified
                     model_fit_kwargs['sample_weight_val'] = w_val.values/w_val.mean()
-            ens_sample_weight =  kwargs.get('ens_sample_weight', None)
+            ens_sample_weight = kwargs.get('ens_sample_weight', None)
             if ens_sample_weight is not None:
                 model_fit_kwargs['sample_weight'] = ens_sample_weight  # sample weights to use for weighted ensemble only
 
@@ -1171,6 +1210,13 @@ def _train_multi_repeats(self, X, y, models: list, n_repeats, n_repeat_start=1,
                     break
             logger.log(20, f'Repeating k-fold bagging: {n+1}/{n_repeats}')
             for i, model in enumerate(models_valid):
+                if not self.get_model_attribute(model=model, attribute='can_fit'):
+                    if isinstance(model, str):
+                        models_valid_next.append(model)
+                    else:
+                        models_valid_next.append(model.name)
+                    continue
+
                 if isinstance(model, str):
                     model = self.load_model(model)
                 if not isinstance(model, BaggedEnsembleModel):
@@ -1304,7 +1350,7 @@ def _train_multi(self, X, y, models: List[AbstractModel], hyperparameter_tune_kw
                 time_limit = time_limit - (time.time() - time_start)
         else:
             model_names_trained = models
-        if (n_repeats > 1) and self.bagged_mode and (n_repeat_start < n_repeats):
+        if (n_repeats > 1) and (n_repeat_start < n_repeats):
             model_names_trained = self._train_multi_repeats(X=X, y=y, models=model_names_trained,
                                                             k_fold=k_fold, n_repeats=n_repeats, n_repeat_start=n_repeat_start, time_limit=time_limit, time_limit_total_level=time_limit_total_level, **kwargs)
         return model_names_trained
@@ -1917,19 +1963,11 @@ def distill(self, X=None, y=None, X_val=None, y_val=None, X_unlabeled=None,
             if y is not None:
                 raise ValueError("X cannot be None when y specified.")
             X = self.load_X()
-            if not self.bagged_mode:
-                try:
-                    X_val = self.load_X_val()
-                except FileNotFoundError:
-                    pass
+            X_val = self.load_X_val()
 
         if y is None:
             y = self.load_y()
-            if not self.bagged_mode:
-                try:
-                    y_val = self.load_y_val()
-                except FileNotFoundError:
-                    pass
+            y_val = self.load_y_val()
 
         if X_val is None:
             if y_val is not None:
diff --git a/tabular/src/autogluon/tabular/trainer/auto_trainer.py b/tabular/src/autogluon/tabular/trainer/auto_trainer.py
index 6b0467bf2664..08c689a8af55 100644
--- a/tabular/src/autogluon/tabular/trainer/auto_trainer.py
+++ b/tabular/src/autogluon/tabular/trainer/auto_trainer.py
@@ -27,25 +27,35 @@ def construct_model_templates(self, hyperparameters, **kwargs):
                                  hyperparameters=hyperparameters, invalid_model_names=invalid_model_names,
                                  silent=silent, **kwargs)
 
-    def fit(self, X, y, hyperparameters, X_val=None, y_val=None, X_unlabeled=None, feature_prune=False, holdout_frac=0.1, num_stack_levels=0, core_kwargs: dict = None, time_limit=None, **kwargs):
+    def fit(self, X, y, hyperparameters, X_val=None, y_val=None, X_unlabeled=None, feature_prune=False, holdout_frac=0.1, num_stack_levels=0, core_kwargs: dict = None, time_limit=None, use_bag_holdout=False, **kwargs):
         for key in kwargs:
             logger.warning(f'Warning: Unknown argument passed to `AutoTrainer.fit()`. Argument: {key}')
 
-        if self.bagged_mode:
-            if (y_val is not None) or (X_val is not None):
-                # TODO: User could be intending to blend instead. Add support for blend stacking.
-                #  This error message is necessary because when calculating out-of-fold predictions for user, we want to return them in the form given in train_data,
-                #  but if we merge train and val here, it becomes very confusing from a users perspective, especially because we reset index, making it impossible to match
-                #  the original train_data to the out-of-fold predictions from `predictor.get_oof_pred_proba()`.
-                raise AssertionError('X_val, y_val is not None, but bagged mode was specified. If calling from `TabularPredictor.fit()`, `tuning_data` must be None.\n'
-                                     'Bagged mode does not use tuning data / validation data. Instead, all data (`train_data` and `tuning_data`) should be combined and specified as `train_data`.\n'
-                                     'Bagging/Stacking with a held-out validation set (blend stacking) is not yet supported.')
-            X_val = None
-            y_val = None
-        else:
-            if (y_val is None) or (X_val is None):
-                X, X_val, y, y_val = generate_train_test_split(X, y, problem_type=self.problem_type, test_size=holdout_frac, random_state=self.random_state)
+        if (y_val is None) or (X_val is None):
+            if not self.bagged_mode or use_bag_holdout:
+                if self.bagged_mode:
+                    # Need at least 2 samples of each class in train data after split for downstream k-fold splits
+                    # to ensure each k-fold has at least 1 sample of each class in training data
+                    min_cls_count_train = 2
+                else:
+                    min_cls_count_train = 1
+                X, X_val, y, y_val = generate_train_test_split(
+                    X,
+                    y,
+                    problem_type=self.problem_type,
+                    test_size=holdout_frac,
+                    random_state=self.random_state,
+                    min_cls_count_train=min_cls_count_train,
+                )
                 logger.log(20, f'Automatically generating train/validation split with holdout_frac={holdout_frac}, Train Rows: {len(X)}, Val Rows: {len(X_val)}')
+        elif self.bagged_mode:
+            # TODO: User could be intending to blend instead. Add support for blend stacking.
+            #  This error message is necessary because when calculating out-of-fold predictions for user, we want to return them in the form given in train_data,
+            #  but if we merge train and val here, it becomes very confusing from a users perspective, especially because we reset index, making it impossible to match
+            #  the original train_data to the out-of-fold predictions from `predictor.get_oof_pred_proba()`.
+            raise AssertionError('X_val, y_val is not None, but bagged mode was specified. If calling from `TabularPredictor.fit()`, `tuning_data` must be None.\n'
+                                 'Bagged mode does not use tuning data / validation data. Instead, all data (`train_data` and `tuning_data`) should be combined and specified as `train_data`.\n'
+                                 'Bagging/Stacking with a held-out validation set (blend stacking) is not yet supported.')
 
         self._train_multi_and_ensemble(X, y, X_val, y_val, X_unlabeled=X_unlabeled, hyperparameters=hyperparameters,
                                        feature_prune=feature_prune,
diff --git a/tabular/src/autogluon/tabular/trainer/model_presets/presets.py b/tabular/src/autogluon/tabular/trainer/model_presets/presets.py
index e6af1c453b01..d13dce0f9924 100644
--- a/tabular/src/autogluon/tabular/trainer/model_presets/presets.py
+++ b/tabular/src/autogluon/tabular/trainer/model_presets/presets.py
@@ -5,7 +5,7 @@
 
 from autogluon.core.metrics import mean_squared_error
 from autogluon.core.constants import AG_ARGS, AG_ARGS_FIT, AG_ARGS_ENSEMBLE, BINARY, MULTICLASS, REGRESSION, SOFTCLASS, QUANTILE
-from autogluon.core.models import AbstractModel, GreedyWeightedEnsembleModel, StackerEnsembleModel
+from autogluon.core.models import AbstractModel, GreedyWeightedEnsembleModel, StackerEnsembleModel, SimpleWeightedEnsembleModel
 
 from .presets_custom import get_preset_custom
 from ..utils import process_hyperparameters
@@ -68,6 +68,7 @@
     AG_IMAGE_NN=ImagePredictorModel,
     FASTTEXT=FastTextModel,
     ENS_WEIGHTED=GreedyWeightedEnsembleModel,
+    SIMPLE_ENS_WEIGHTED=SimpleWeightedEnsembleModel,
 )
 
 DEFAULT_MODEL_NAMES = {
@@ -86,6 +87,7 @@
     ImagePredictorModel: 'ImagePredictor',
     FastTextModel: 'FastText',
     GreedyWeightedEnsembleModel: 'WeightedEnsemble',
+    SimpleWeightedEnsembleModel: 'WeightedEnsemble',
 }
 
 
diff --git a/tabular/tests/unittests/test_tabular.py b/tabular/tests/unittests/test_tabular.py
index 60a0e76d90e4..0a5aa4f039e0 100644
--- a/tabular/tests/unittests/test_tabular.py
+++ b/tabular/tests/unittests/test_tabular.py
@@ -203,7 +203,7 @@ def run_tabular_benchmark_toy(fit_args):
         raise AssertionError(f'{dataset["name"]} should raise an exception.')
 
 
-def run_tabular_benchmarks(fast_benchmark, subsample_size, perf_threshold, seed_val, fit_args, dataset_indices=None, run_distill=False):
+def run_tabular_benchmarks(fast_benchmark, subsample_size, perf_threshold, seed_val, fit_args, dataset_indices=None, run_distill=False, crash_in_oof=False):
     print("Running fit with args:")
     print(fit_args)
     # Each train/test dataset must be located in single directory with the given names.
@@ -289,7 +289,7 @@ def run_tabular_benchmarks(fast_benchmark, subsample_size, perf_threshold, seed_
             if (not fast_benchmark) and (performance_vals[idx] > dataset['performance_val'] * perf_threshold):
                 warnings.warn("Performance on dataset %s is %s times worse than previous performance." %
                               (dataset['name'], performance_vals[idx]/(EPS+dataset['performance_val'])))
-            if predictor._trainer.bagged_mode:
+            if predictor._trainer.bagged_mode and not crash_in_oof:
                 # TODO: Test index alignment with original training data (first handle duplicated rows / dropped rows edge cases)
                 y_pred_oof = predictor.get_oof_pred()
                 y_pred_proba_oof = predictor.get_oof_pred_proba(as_multiclass=False)
@@ -664,3 +664,45 @@ def test_tabular_bagstack():
     run_tabular_benchmarks(fast_benchmark=fast_benchmark, subsample_size=subsample_size, perf_threshold=perf_threshold,
                            seed_val=seed_val, fit_args=fit_args, run_distill=True)
 
+
+@pytest.mark.slow
+def test_tabular_bagstack_use_bag_holdout():
+    ############ Benchmark options you can set: ########################
+    num_stack_levels = 2
+    num_bag_folds = 3
+    perf_threshold = 1.1  # How much worse can performance on each dataset be vs previous performance without warning
+    seed_val = 53  # random seed
+    subsample_size = None
+    hyperparameter_tune_kwargs = None
+    verbosity = 2  # how much output to print
+    hyperparameters = None
+    time_limit = None
+    fast_benchmark = True  # False
+    # If True, run a faster benchmark (subsample training sets, less epochs, etc),
+    # otherwise we run full benchmark with default AutoGluon settings.
+    # performance_value warnings are disabled when fast_benchmark = True.
+
+    #### If fast_benchmark = True, can control model training time here. Only used if fast_benchmark=True ####
+    if fast_benchmark:
+        subsample_size = 105
+        nn_options = {'num_epochs': 2}
+        gbm_options = [{'num_boost_round': 40}, 'GBMLarge']
+        hyperparameters = {'GBM': gbm_options, 'NN': nn_options}
+        time_limit = 60
+
+    fit_args = {
+        'num_bag_folds': num_bag_folds,
+        'num_stack_levels': num_stack_levels,
+        'verbosity': verbosity,
+        'use_bag_holdout': True,
+    }
+    if hyperparameter_tune_kwargs is not None:
+        fit_args['hyperparameter_tune_kwargs'] = hyperparameter_tune_kwargs
+    if hyperparameters is not None:
+        fit_args['hyperparameters'] = hyperparameters
+    if time_limit is not None:
+        fit_args['time_limit'] = time_limit
+        fit_args['num_bag_sets'] = 2
+    ###################################################################
+    run_tabular_benchmarks(fast_benchmark=fast_benchmark, subsample_size=subsample_size, perf_threshold=perf_threshold,
+                           seed_val=seed_val, fit_args=fit_args, run_distill=True, crash_in_oof=True)


Code review comments:
diff_hunk: @@ -240,13 +259,21 @@ def augment_rare_classes(X, label, threshold):
             new_df = new_df.append(clss_df.copy())
         aug_df = aug_df.append(new_df.copy())
 
+    # Ensure new samples generated via augmentation have unique indices
+    aug_df = aug_df.reset_index(drop=True)
+    aug_df_len = len(aug_df)
+    X_index_aug_start = max(list(X.index)) + 1
path: core/src/autogluon/core/utils/utils.py
text: @user1:
```python
X_index_aug_start = X.index.max() + 1
``` 

@author:
Added

diff_hunk: @@ -47,7 +49,24 @@ def generate_kfold(X, y=None, n_splits=5, random_state=0, stratified=False, n_re
             kf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
 
         kf.get_n_splits(X, y)
-        return [[train_index, test_index] for train_index, test_index in kf.split(X, y)]
+        # FIXME: There is a bug in sklearn that causes an incorrect ValueError if performing stratification and all classes have fewer than n_splits samples.
+        #  This is hacked by adding a dummy class with n_splits samples, performing the kfold split, then removing the dummy samples from all resulting indices.
+        #  This is very inefficient and complicated and ideally should be fixed in sklearn.
+        with warning_filter():
+            try:
+                out = [[train_index, test_index] for train_index, test_index in kf.split(X, y)]
+            except:
+                y_dummy = pd.concat([y, pd.Series([-1] * n_splits)], ignore_index=True)
+                X_dummy = pd.concat([X, X.head(n_splits)], ignore_index=True)
+                invalid_index = set(list(y_dummy.tail(n_splits).index))
+                out = [[train_index, test_index] for train_index, test_index in kf.split(X_dummy, y_dummy)]
+                len_out = len(out)
+                for i in range(len_out):
+                    train_index = out[i][0]
+                    test_index = out[i][1]
path: core/src/autogluon/core/utils/utils.py
text: @user1:
nit:
```python
train_index, test_index = out[i]
```

@author:
Added