implement optimizer state with opt offloading (#4141)

linefb · facebook-github-bot · commit 16ce7727ecc0 · 2025-05-21T16:06:19.000-07:00
Summary: Pull Request resolved: #4141 X-link: facebookresearch/FBGEMM#1224 implement split_optimizer_states for optimizer state dict integration Differential Revision: D74790121 Reviewed By: duduyi2013, bobbyliujb
diff --git a/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py b/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py
@@ -2007,6 +2007,99 @@ def split_optimizer_states(
         )
         return opt_list
 
+        opt_list = []
+        table_offset = 0
+
+        dtype = self.weights_precision.as_dtype()
+        optimizer_dim = self.optimizer.state_size_dim(dtype)
+        pad4_optimizer_dim = pad4(optimizer_dim)
+        logging.info(
+            f"split_optimizer_states: {optimizer_dim=} {pad4_optimizer_dim=} {self.optimizer.dtype()=} {self.enable_load_state_dict_mode=}"
+        )
+
+        for t, (emb_height, emb_dim) in enumerate(self.embedding_specs):
+            # pyre-ignore
+            bucket_id_start, _ = self.kv_zch_params.bucket_offsets[t]
+            # pyre-ignore
+            bucket_size = self.kv_zch_params.bucket_sizes[t]
+            row_offset = table_offset
+            if sorted_id_tensor is None or sorted_id_tensor[t].numel() == 0:
+                opt_list.append(
+                    torch.empty(0, dtype=self.optimizer.dtype(), device="cpu")
+                    # empty optimizer state for module initialization
+                )
+            else:
+                if not self.enable_optimizer_offloading:
+                    # convert global id back to local id, then linearize with table offset
+                    local_id_tensor = (
+                        sorted_id_tensor[t]
+                        - bucket_id_start * bucket_size
+                        + table_offset
+                    )
+                    opt_list.append(
+                        self.momentum1_dev.detach().cpu()[local_id_tensor].view(-1),
+                    )
+                else:
+                    emb_opt_dim = pad4(emb_dim) + pad4_optimizer_dim
+                    row_offset = table_offset - (bucket_id_start * bucket_size)
+                    # using KVTensorWrapper to query backend to avoid OOM memory, since
+                    # backend will return both weight and optimizer in one tensor, read the whole tensor
+                    # out could OOM CPU memory.
+                    tensor_wrapper = torch.classes.fbgemm.KVTensorWrapper(
+                        shape=[emb_height, emb_opt_dim],
+                        dtype=dtype,
+                        row_offset=row_offset,
+                        snapshot_handle=snapshot_handle,
+                        materialized_shape=([sorted_id_tensor[t].size(0), emb_opt_dim]),
+                        sorted_indices=sorted_id_tensor[t],
+                    )
+                    (
+                        tensor_wrapper.set_embedding_rocks_dp_wrapper(self.ssd_db)
+                        if self.backend_type == BackendType.SSD
+                        else tensor_wrapper.set_dram_db_wrapper(self.ssd_db)
+                    )
+                    opt_list.append(
+                        self.get_offloaded_optimizer_states(
+                            tensor_wrapper=tensor_wrapper,
+                            row=sorted_id_tensor[t].size(
+                                0
+                            ),  # we only need to copy the size of sorted_id_tensor
+                            optimizer_dim=optimizer_dim,
+                            start_dim_pos=pad4(emb_dim),
+                        )
+                    )
+            table_offset += emb_height
+        logging.info(
+            f"KV ZCH tables split_optimizer_states query latency: {(time.time() - start_time) * 1000} ms"
+        )
+        return opt_list
+
+    @torch.jit.export
+    def get_offloaded_optimizer_states(
+        self,
+        # pyre-ignore [2]
+        tensor_wrapper,
+        row: int,
+        optimizer_dim: int,
+        start_dim_pos: int,
+    ) -> torch.Tensor:
+        weight_dtype = self.weights_precision.as_dtype()
+        opt_state_t = torch.empty(
+            row, optimizer_dim, dtype=weight_dtype, device="cpu"
+        )  # 1D optimizer for OptimType.EXACT_ROWWISE_ADAGRAD
+
+        # pyre-ignore [16]
+        chunk_size = self.kv_zch_params.streaming_ckpt_chunk_size
+        for i in range(0, row, chunk_size):
+            length = min(chunk_size, row - i)
+            opt_state_t.narrow(0, i, length).copy_(
+                tensor_wrapper.narrow(0, i, length).narrow(
+                    1, start_dim_pos, optimizer_dim
+                )
+            )
+        # view optimizer state back to correct dtype
+        return opt_state_t.view(-1).view(self.optimizer.dtype())
+
     @torch.jit.export
     def get_optimizer_state(
         self,
diff --git a/fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py b/fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py
@@ -2007,3 +2007,230 @@ def test_kv_opt_state_w_offloading(
                 atol=tolerance,
                 rtol=tolerance,
             )
+
+    @given(
+        **default_st,
+        num_buckets=st.integers(min_value=10, max_value=15),
+    )
+    @settings(verbosity=Verbosity.verbose, max_examples=MAX_EXAMPLES, deadline=None)
+    def test_kv_opt_state_w_offloading(
+        self,
+        T: int,
+        D: int,
+        B: int,
+        log_E: int,
+        L: int,
+        weighted: bool,
+        cache_set_scale: float,
+        pooling_mode: PoolingMode,
+        weights_precision: SparseType,
+        output_dtype: SparseType,
+        share_table: bool,
+        trigger_bounds_check: bool,
+        mixed_B: bool,
+        num_buckets: int,
+    ) -> None:
+        # Constants
+        lr = 0.5
+        eps = 0.2
+        ssd_shards = 2
+
+        trigger_bounds_check = False  # don't stimulate boundary check cases
+        assume(not weighted or pooling_mode == PoolingMode.SUM)
+        assume(not mixed_B or pooling_mode != PoolingMode.NONE)
+
+        # TODO: check split_optimizer_states when optimizer offloading is ready
+        # Generate embedding modules and inputs
+        (
+            emb,
+            emb_ref,
+            Es,
+            _,
+            bucket_offsets,
+            bucket_sizes,
+        ) = self.generate_kvzch_tbes(
+            T,
+            D,
+            B,
+            log_E,
+            L,
+            weighted,
+            lr=lr,
+            eps=eps,
+            ssd_shards=ssd_shards,
+            cache_set_scale=cache_set_scale,
+            pooling_mode=pooling_mode,
+            weights_precision=weights_precision,
+            output_dtype=output_dtype,
+            share_table=share_table,
+            num_buckets=num_buckets,
+            enable_optimizer_offloading=False,
+        )
+
+        # Generate inputs
+        (
+            indices_list,
+            per_sample_weights_list,
+            indices,
+            offsets,
+            per_sample_weights,
+            batch_size_per_feature_per_rank,
+        ) = self.generate_inputs_(
+            B,
+            L,
+            Es,
+            emb.feature_table_map,
+            weights_precision=weights_precision,
+            trigger_bounds_check=trigger_bounds_check,
+            mixed_B=mixed_B,
+            bucket_offsets=bucket_offsets,
+            bucket_sizes=bucket_sizes,
+            is_kv_tbes=True,
+        )
+
+        # Execute forward
+        output_ref_list, output = self.execute_ssd_forward_(
+            emb,
+            emb_ref,
+            indices_list,
+            per_sample_weights_list,
+            indices,
+            offsets,
+            per_sample_weights,
+            B,
+            L,
+            weighted,
+            batch_size_per_feature_per_rank=batch_size_per_feature_per_rank,
+        )
+
+        # Generate output gradient
+        output_grad_list = [torch.randn_like(out) for out in output_ref_list]
+
+        # Execute torch EmbeddingBag backward
+        [out.backward(grad) for (out, grad) in zip(output_ref_list, output_grad_list)]
+        if batch_size_per_feature_per_rank is not None:
+            grad_test = self.concat_ref_tensors_vbe(
+                output_grad_list, batch_size_per_feature_per_rank
+            )
+        else:
+            grad_test = self.concat_ref_tensors(
+                output_grad_list,
+                pooling_mode != PoolingMode.NONE,  # do_pooling
+                B,
+                D * 4,
+            )
+
+        # Execute TBE SSD backward
+        output.backward(grad_test)
+
+        tolerance = (
+            1.0e-4
+            if weights_precision == SparseType.FP32 and output_dtype == SparseType.FP32
+            else 1.0e-2
+        )
+
+        emb.flush()
+
+        # Compare emb state dict with expected values from nn.EmbeddingBag
+        emb_state_dict_list, bucket_asc_ids_list, num_active_id_per_bucket_list = (
+            emb.split_embedding_weights(no_snapshot=False, should_flush=True)
+        )
+        split_optimizer_states = emb.split_optimizer_states(bucket_asc_ids_list)
+        table_input_id_range = []
+        for t, row in enumerate(Es):
+            bucket_id_start = bucket_offsets[t][0]
+            bucket_id_end = bucket_offsets[t][1]
+            bucket_size = bucket_sizes[t]
+            table_input_id_range.append(
+                (
+                    min(bucket_id_start * bucket_size, row),
+                    min(bucket_id_end * bucket_size, row),
+                )
+            )
+            # since we use ref_emb in dense format, the rows start from id 0
+            self.assertEqual(table_input_id_range[-1][0], 0)
+
+        # Compare optimizer states
+        for f, t in self.get_physical_table_arg_indices_(emb.feature_table_map):
+            # pyre-fixme[16]: Optional type has no attribute `float`.
+            ref_emb = emb_ref[f].weight.grad.float().to_dense().pow(2).cpu()
+            ref_optimizer_state = ref_emb.mean(dim=1)[
+                table_input_id_range[t][0] : min(
+                    table_input_id_range[t][1], emb_ref[f].weight.size(0)
+                )
+            ]
+            ref_kv_opt = ref_optimizer_state[bucket_asc_ids_list[t]].view(-1)
+            torch.testing.assert_close(
+                split_optimizer_states[t].float(),
+                ref_kv_opt,
+                atol=tolerance,
+                rtol=tolerance,
+            )
+
+        for feature_index, table_index in self.get_physical_table_arg_indices_(
+            emb.feature_table_map
+        ):
+            """
+            validate bucket_asc_ids_list and num_active_id_per_bucket_list
+            """
+            bucket_asc_id = bucket_asc_ids_list[table_index]
+            num_active_id_per_bucket = num_active_id_per_bucket_list[table_index]
+
+            bucket_id_start = bucket_offsets[table_index][0]
+            bucket_id_offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(
+                num_active_id_per_bucket.view(-1)
+            )
+            for bucket_idx, id_count in enumerate(num_active_id_per_bucket):
+                bucket_id = bucket_idx + bucket_id_start
+                active_id_cnt = 0
+                for idx in range(
+                    bucket_id_offsets[bucket_idx],
+                    bucket_id_offsets[bucket_idx + 1],
+                ):
+                    # for chunk-based hashing
+                    self.assertEqual(
+                        bucket_id, bucket_asc_id[idx] // bucket_sizes[table_index]
+                    )
+                    active_id_cnt += 1
+                self.assertEqual(active_id_cnt, id_count)
+
+            """
+            validate embeddings
+            """
+            num_ids = len(bucket_asc_ids_list[table_index])
+            emb_r_w = emb_ref[feature_index].weight[
+                bucket_asc_ids_list[table_index].view(-1)
+            ]
+            emb_r_w_g = (
+                emb_ref[feature_index]
+                .weight.grad.float()
+                .to_dense()[bucket_asc_ids_list[table_index].view(-1)]
+            )
+            self.assertLess(table_index, len(emb_state_dict_list))
+            assert len(split_optimizer_states[table_index]) == num_ids
+            opt = split_optimizer_states[table_index]
+            new_ref_weight = torch.addcdiv(
+                emb_r_w.float(),
+                value=-lr,
+                tensor1=emb_r_w_g,
+                tensor2=opt.float()
+                .sqrt_()
+                .add_(eps)
+                .view(
+                    num_ids,
+                    1,
+                )
+                .cuda(),
+            ).cpu()
+
+            emb_w = (
+                emb_state_dict_list[table_index]
+                .narrow(0, 0, bucket_asc_ids_list[table_index].size(0))
+                .float()
+            )
+            torch.testing.assert_close(
+                emb_w,
+                new_ref_weight,
+                atol=tolerance,
+                rtol=tolerance,
+            )
