diff --git a/drivers/gpu/nvgpu/common/fifo/tsg.c b/drivers/gpu/nvgpu/common/fifo/tsg.c
index e4070c26f..ae2ad7364 100644
--- a/drivers/gpu/nvgpu/common/fifo/tsg.c
+++ b/drivers/gpu/nvgpu/common/fifo/tsg.c
@@ -337,14 +337,10 @@ void nvgpu_tsg_cleanup_sw(struct gk20a *g)
 	nvgpu_mutex_destroy(&f->tsg_inuse_mutex);
 }
 
-int nvgpu_tsg_init_support(struct gk20a *g, u32 tsgid)
+static void nvgpu_tsg_init_support(struct gk20a *g, u32 tsgid)
 {
 	struct nvgpu_tsg *tsg = NULL;
 
-	if (tsgid >= g->fifo.num_channels) {
-		return -EINVAL;
-	}
-
 	tsg = &g->fifo.tsg[tsgid];
 
 	tsg->in_use = false;
@@ -357,13 +353,12 @@ int nvgpu_tsg_init_support(struct gk20a *g, u32 tsgid)
 	nvgpu_init_list_node(&tsg->event_id_list);
 
 	nvgpu_mutex_init(&tsg->event_id_list_lock);
-	return 0;
 }
 
 int nvgpu_tsg_setup_sw(struct gk20a *g)
 {
 	struct nvgpu_fifo *f = &g->fifo;
-	u32 tsgid, i;
+	u32 tsgid;
 	int err;
 
 	nvgpu_mutex_init(&f->tsg_inuse_mutex);
@@ -376,24 +371,11 @@ int nvgpu_tsg_setup_sw(struct gk20a *g)
 	}
 
 	for (tsgid = 0; tsgid < f->num_channels; tsgid++) {
-		err = nvgpu_tsg_init_support(g, tsgid);
-		if (err != 0) {
-			nvgpu_err(g, "tsg init failed, tsgid=%u", tsgid);
-			goto clean_up;
-		}
+		nvgpu_tsg_init_support(g, tsgid);
 	}
 
 	return 0;
 
-clean_up:
-	for (i = 0; i < tsgid; i++) {
-		struct nvgpu_tsg *tsg = &g->fifo.tsg[i];
-
-		nvgpu_tsg_destroy(g, tsg);
-	}
-	nvgpu_vfree(g, f->tsg);
-	f->tsg = NULL;
-
 clean_up_mutex:
 	nvgpu_mutex_destroy(&f->tsg_inuse_mutex);
 	return err;
diff --git a/drivers/gpu/nvgpu/include/nvgpu/tsg.h b/drivers/gpu/nvgpu/include/nvgpu/tsg.h
index b92bdb9b9..0956ea7be 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/tsg.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/tsg.h
@@ -99,7 +99,6 @@ struct nvgpu_tsg *nvgpu_tsg_open(struct gk20a *g, pid_t pid);
 void nvgpu_tsg_release_common(struct gk20a *g, struct nvgpu_tsg *tsg);
 void nvgpu_tsg_release(struct nvgpu_ref *ref);
 
-int nvgpu_tsg_init_support(struct gk20a *g, u32 tsgid);
 int nvgpu_tsg_setup_sw(struct gk20a *g);
 void nvgpu_tsg_cleanup_sw(struct gk20a *g);
 
diff --git a/drivers/gpu/nvgpu/libnvgpu-drv_safe.export b/drivers/gpu/nvgpu/libnvgpu-drv_safe.export
index f3588e953..da525fccb 100644
--- a/drivers/gpu/nvgpu/libnvgpu-drv_safe.export
+++ b/drivers/gpu/nvgpu/libnvgpu-drv_safe.export
@@ -166,9 +166,11 @@ nvgpu_rwsem_init
 nvgpu_tsg_abort
 nvgpu_tsg_bind_channel
 nvgpu_tsg_check_and_get_from_id
+nvgpu_tsg_cleanup_sw
 nvgpu_tsg_default_timeslice_us
 nvgpu_tsg_open
 nvgpu_tsg_release
+nvgpu_tsg_setup_sw
 nvgpu_tsg_unbind_channel
 nvgpu_tsg_unbind_channel_check_hw_state
 nvgpu_tsg_unbind_channel_check_ctx_reload
diff --git a/userspace/units/fifo/tsg/nvgpu-tsg.c b/userspace/units/fifo/tsg/nvgpu-tsg.c
index 69695e7ba..67f832eac 100644
--- a/userspace/units/fifo/tsg/nvgpu-tsg.c
+++ b/userspace/units/fifo/tsg/nvgpu-tsg.c
@@ -37,6 +37,7 @@
 #include <nvgpu/gr/ctx.h>
 
 #include "common/gr/ctx_priv.h"
+#include <nvgpu/posix/posix-fault-injection.h>
 
 #include "hal/fifo/tsg_gk20a.h"
 
@@ -1135,7 +1136,64 @@ done:
 	return rc;
 }
 
+#define F_TSG_SETUP_SW_VZALLOC_FAIL		BIT(0)
+#define F_TSG_SETUP_SW_LAST			BIT(1)
+
+static const char *f_tsg_setup_sw[] = {
+	"vzalloc_fail",
+};
+
+static int test_tsg_setup_sw(struct unit_module *m,
+		struct gk20a *g, void *args)
+{
+	struct gpu_ops gops = g->ops;
+	struct nvgpu_posix_fault_inj *kmem_fi;
+	u32 branches = 0U;
+	int rc = UNIT_FAIL;
+	int err;
+	u32 fail = F_TSG_SETUP_SW_VZALLOC_FAIL;
+	u32 prune = fail;
+
+	kmem_fi = nvgpu_kmem_get_fault_injection();
+
+	for (branches = 0U; branches < F_TSG_SETUP_SW_LAST; branches++) {
+
+		if (pruned(branches, prune)) {
+			unit_verbose(m, "%s branches=%s (pruned)\n", __func__,
+				branches_str(branches, f_tsg_setup_sw));
+			continue;
+		}
+		subtest_setup(branches);
+		unit_verbose(m, "%s branches=%s\n", __func__,
+			branches_str(branches, f_tsg_setup_sw));
+
+		nvgpu_posix_enable_fault_injection(kmem_fi,
+			branches & F_TSG_SETUP_SW_VZALLOC_FAIL ?
+			true : false, 0);
+
+		err = nvgpu_tsg_setup_sw(g);
+
+		if (branches & fail) {
+			assert(err != 0);
+		} else {
+			assert(err == 0);
+			nvgpu_tsg_cleanup_sw(g);
+		}
+	}
+
+	rc = UNIT_SUCCESS;
+done:
+	nvgpu_posix_enable_fault_injection(kmem_fi, false, 0);
+	if (rc != UNIT_SUCCESS) {
+		unit_err(m, "%s branches=%s\n", __func__,
+			branches_str(branches, f_tsg_setup_sw));
+	}
+	g->ops = gops;
+	return rc;
+}
+
 struct unit_module_test nvgpu_tsg_tests[] = {
+	UNIT_TEST(setup_sw, test_tsg_setup_sw, &unit_ctx, 0),
 	UNIT_TEST(init_support, test_fifo_init_support, &unit_ctx, 0),
 	UNIT_TEST(open, test_tsg_open, &unit_ctx, 0),
 	UNIT_TEST(release, test_tsg_release, &unit_ctx, 0),