diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
index 6c31306274..73713e0ff8 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
@@ -127,11 +127,11 @@ export function* run(
     code,
     useMemoCacheIdentifier,
   );
-  yield {
+  yield log({
     kind: 'debug',
     name: 'EnvironmentConfig',
     value: prettyFormat(env.config),
-  };
+  });
   const ast = yield* runWithEnvironment(func, env);
   return ast;
 }
diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts
index aae7b82594..2d9e21af1d 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts
@@ -335,6 +335,7 @@ function extractManualMemoizationArgs(
 export function dropManualMemoization(func: HIRFunction): void {
   const isValidationEnabled =
     func.env.config.validatePreserveExistingMemoizationGuarantees ||
+    func.env.config.validateNoSetStateInRender ||
     func.env.config.enablePreserveExistingMemoizationGuarantees;
   const sidemap: IdentifierSidemap = {
     functions: new Map(),
diff --git a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoSetStateInRender.ts b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoSetStateInRender.ts
index 5c6339a2be..3f378b1289 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoSetStateInRender.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoSetStateInRender.ts
@@ -6,7 +6,7 @@
  */
 
 import {CompilerError, ErrorSeverity} from '../CompilerError';
-import {HIRFunction, IdentifierId, Place, isSetStateType} from '../HIR';
+import {HIRFunction, IdentifierId, isSetStateType} from '../HIR';
 import {computeUnconditionalBlocks} from '../HIR/ComputeUnconditionalBlocks';
 import {eachInstructionValueOperand} from '../HIR/visitors';
 import {Err, Ok, Result} from '../Utils/Result';
@@ -49,63 +49,97 @@ function validateNoSetStateInRenderImpl(
   unconditionalSetStateFunctions: Set<IdentifierId>,
 ): Result<void, CompilerError> {
   const unconditionalBlocks = computeUnconditionalBlocks(fn);
-
+  let activeManualMemoId: number | null = null;
   const errors = new CompilerError();
   for (const [, block] of fn.body.blocks) {
-    if (unconditionalBlocks.has(block.id)) {
-      for (const instr of block.instructions) {
-        switch (instr.value.kind) {
-          case 'LoadLocal': {
-            if (
-              unconditionalSetStateFunctions.has(
-                instr.value.place.identifier.id,
-              )
-            ) {
-              unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
-            }
-            break;
+    for (const instr of block.instructions) {
+      switch (instr.value.kind) {
+        case 'LoadLocal': {
+          if (
+            unconditionalSetStateFunctions.has(instr.value.place.identifier.id)
+          ) {
+            unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
           }
-          case 'StoreLocal': {
-            if (
-              unconditionalSetStateFunctions.has(
-                instr.value.value.identifier.id,
-              )
-            ) {
-              unconditionalSetStateFunctions.add(
-                instr.value.lvalue.place.identifier.id,
-              );
-              unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
-            }
-            break;
-          }
-          case 'ObjectMethod':
-          case 'FunctionExpression': {
-            if (
-              // faster-path to check if the function expression references a setState
-              [...eachInstructionValueOperand(instr.value)].some(
-                operand =>
-                  isSetStateType(operand.identifier) ||
-                  unconditionalSetStateFunctions.has(operand.identifier.id),
-              ) &&
-              // if yes, does it unconditonally call it?
-              validateNoSetStateInRenderImpl(
-                instr.value.loweredFunc.func,
-                unconditionalSetStateFunctions,
-              ).isErr()
-            ) {
-              // This function expression unconditionally calls a setState
-              unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
-            }
-            break;
-          }
-          case 'CallExpression': {
-            validateNonSetState(
-              errors,
-              unconditionalSetStateFunctions,
-              instr.value.callee,
+          break;
+        }
+        case 'StoreLocal': {
+          if (
+            unconditionalSetStateFunctions.has(instr.value.value.identifier.id)
+          ) {
+            unconditionalSetStateFunctions.add(
+              instr.value.lvalue.place.identifier.id,
             );
-            break;
+            unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
           }
+          break;
+        }
+        case 'ObjectMethod':
+        case 'FunctionExpression': {
+          if (
+            // faster-path to check if the function expression references a setState
+            [...eachInstructionValueOperand(instr.value)].some(
+              operand =>
+                isSetStateType(operand.identifier) ||
+                unconditionalSetStateFunctions.has(operand.identifier.id),
+            ) &&
+            // if yes, does it unconditonally call it?
+            validateNoSetStateInRenderImpl(
+              instr.value.loweredFunc.func,
+              unconditionalSetStateFunctions,
+            ).isErr()
+          ) {
+            // This function expression unconditionally calls a setState
+            unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
+          }
+          break;
+        }
+        case 'StartMemoize': {
+          CompilerError.invariant(activeManualMemoId === null, {
+            reason: 'Unexpected nested StartMemoize instructions',
+            loc: instr.value.loc,
+          });
+          activeManualMemoId = instr.value.manualMemoId;
+          break;
+        }
+        case 'FinishMemoize': {
+          CompilerError.invariant(
+            activeManualMemoId === instr.value.manualMemoId,
+            {
+              reason:
+                'Expected FinishMemoize to align with previous StartMemoize instruction',
+              loc: instr.value.loc,
+            },
+          );
+          activeManualMemoId = null;
+          break;
+        }
+        case 'CallExpression': {
+          const callee = instr.value.callee;
+          if (
+            isSetStateType(callee.identifier) ||
+            unconditionalSetStateFunctions.has(callee.identifier.id)
+          ) {
+            if (activeManualMemoId !== null) {
+              errors.push({
+                reason:
+                  'Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState)',
+                description: null,
+                severity: ErrorSeverity.InvalidReact,
+                loc: callee.loc,
+                suggestions: null,
+              });
+            } else if (unconditionalBlocks.has(block.id)) {
+              errors.push({
+                reason:
+                  'This is an unconditional set state during render, which will trigger an infinite loop. (https://react.dev/reference/react/useState)',
+                description: null,
+                severity: ErrorSeverity.InvalidReact,
+                loc: callee.loc,
+                suggestions: null,
+              });
+            }
+          }
+          break;
         }
       }
     }
@@ -117,23 +151,3 @@ function validateNoSetStateInRenderImpl(
     return Ok(undefined);
   }
 }
-
-function validateNonSetState(
-  errors: CompilerError,
-  unconditionalSetStateFunctions: Set<IdentifierId>,
-  operand: Place,
-): void {
-  if (
-    isSetStateType(operand.identifier) ||
-    unconditionalSetStateFunctions.has(operand.identifier.id)
-  ) {
-    errors.push({
-      reason:
-        'This is an unconditional set state during render, which will trigger an infinite loop. (https://react.dev/reference/react/useState)',
-      description: null,
-      severity: ErrorSeverity.InvalidReact,
-      loc: typeof operand.loc !== 'symbol' ? operand.loc : null,
-      suggestions: null,
-    });
-  }
-}
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-global-mutation-unused-usecallback.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-global-mutation-unused-usecallback.expect.md
index 68bc0d5913..f5c393c174 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-global-mutation-unused-usecallback.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-global-mutation-unused-usecallback.expect.md
@@ -36,6 +36,9 @@ function Component() {
   }
   return t0;
 }
+function _temp() {
+  window.foo = true;
+}
 
 export const FIXTURE_ENTRYPOINT = {
   fn: Component,
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.expect.md
index 207c884aeb..e8ed54ef26 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.expect.md
@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @validateRefAccessDuringRender
+// @validateRefAccessDuringRender @validateNoSetStateInRender:false
 import {useCallback, useEffect, useRef, useState} from 'react';
 
 function Component() {
@@ -42,7 +42,7 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @validateRefAccessDuringRender
+import { c as _c } from "react/compiler-runtime"; // @validateRefAccessDuringRender @validateNoSetStateInRender:false
 import { useCallback, useEffect, useRef, useState } from "react";
 
 function Component() {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.js
index cf311d4df6..6942904902 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.js
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.js
@@ -1,4 +1,4 @@
-// @validateRefAccessDuringRender
+// @validateRefAccessDuringRender @validateNoSetStateInRender:false
 import {useCallback, useEffect, useRef, useState} from 'react';
 
 function Component() {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-setState-in-useMemo.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-setState-in-useMemo.expect.md
new file mode 100644
index 0000000000..f1666cc401
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-setState-in-useMemo.expect.md
@@ -0,0 +1,36 @@
+
+## Input
+
+```javascript
+function Component({item, cond}) {
+  const [prevItem, setPrevItem] = useState(item);
+  const [state, setState] = useState(0);
+
+  useMemo(() => {
+    if (cond) {
+      setPrevItem(item);
+      setState(0);
+    }
+  }, [cond, key, init]);
+
+  return state;
+}
+
+```
+
+
+## Error
+
+```
+   5 |   useMemo(() => {
+   6 |     if (cond) {
+>  7 |       setPrevItem(item);
+     |       ^^^^^^^^^^^ InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (7:7)
+
+InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (8:8)
+   8 |       setState(0);
+   9 |     }
+  10 |   }, [cond, key, init]);
+```
+          
+      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-setState-in-useMemo.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-setState-in-useMemo.js
new file mode 100644
index 0000000000..4385ef6a6c
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-setState-in-useMemo.js
@@ -0,0 +1,13 @@
+function Component({item, cond}) {
+  const [prevItem, setPrevItem] = useState(item);
+  const [state, setState] = useState(0);
+
+  useMemo(() => {
+    if (cond) {
+      setPrevItem(item);
+      setState(0);
+    }
+  }, [cond, key, init]);
+
+  return state;
+}
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.expect.md
new file mode 100644
index 0000000000..acf84652b2
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.expect.md
@@ -0,0 +1,38 @@
+
+## Input
+
+```javascript
+import {useCallback} from 'react';
+
+function useKeyedState({key, init}) {
+  const [prevKey, setPrevKey] = useState(key);
+  const [state, setState] = useState(init);
+
+  const fn = useCallback(() => {
+    setPrevKey(key);
+    setState(init);
+  });
+
+  useMemo(() => {
+    fn();
+  }, [key, init]);
+
+  return state;
+}
+
+```
+
+
+## Error
+
+```
+  11 |
+  12 |   useMemo(() => {
+> 13 |     fn();
+     |     ^^ InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (13:13)
+  14 |   }, [key, init]);
+  15 |
+  16 |   return state;
+```
+          
+      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.js
new file mode 100644
index 0000000000..dbdf5f5a4d
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.js
@@ -0,0 +1,17 @@
+import {useCallback} from 'react';
+
+function useKeyedState({key, init}) {
+  const [prevKey, setPrevKey] = useState(key);
+  const [state, setState] = useState(init);
+
+  const fn = useCallback(() => {
+    setPrevKey(key);
+    setState(init);
+  });
+
+  useMemo(() => {
+    fn();
+  }, [key, init]);
+
+  return state;
+}
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo.expect.md
new file mode 100644
index 0000000000..aee09b1790
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo.expect.md
@@ -0,0 +1,34 @@
+
+## Input
+
+```javascript
+function useKeyedState({key, init}) {
+  const [prevKey, setPrevKey] = useState(key);
+  const [state, setState] = useState(init);
+
+  useMemo(() => {
+    setPrevKey(key);
+    setState(init);
+  }, [key, init]);
+
+  return state;
+}
+
+```
+
+
+## Error
+
+```
+  4 |
+  5 |   useMemo(() => {
+> 6 |     setPrevKey(key);
+    |     ^^^^^^^^^^ InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (6:6)
+
+InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (7:7)
+  7 |     setState(init);
+  8 |   }, [key, init]);
+  9 |
+```
+          
+      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo.js
new file mode 100644
index 0000000000..ebbb6b2d14
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo.js
@@ -0,0 +1,11 @@
+function useKeyedState({key, init}) {
+  const [prevKey, setPrevKey] = useState(key);
+  const [state, setState] = useState(init);
+
+  useMemo(() => {
+    setPrevKey(key);
+    setState(init);
+  }, [key, init]);
+
+  return state;
+}
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro.expect.md
index dcd252ab58..2e9daceed7 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro.expect.md
@@ -24,12 +24,14 @@ function Component(props) {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function Component(props) {
-  const $ = _c(2);
+  const $ = _c(7);
   const item = props.item;
   let t0;
+  let baseVideos;
+  let thumbnails;
   if ($[0] !== item) {
-    const thumbnails = [];
-    const baseVideos = getBaseVideos(item);
+    thumbnails = [];
+    baseVideos = getBaseVideos(item);
 
     baseVideos.forEach((video) => {
       const baseVideo = video.hasBaseVideo;
@@ -37,14 +39,26 @@ function Component(props) {
         thumbnails.push({ extraVideo: true });
       }
     });
-
-    t0 = <FlatList baseVideos={baseVideos} items={thumbnails} />;
     $[0] = item;
     $[1] = t0;
+    $[2] = baseVideos;
+    $[3] = thumbnails;
   } else {
     t0 = $[1];
+    baseVideos = $[2];
+    thumbnails = $[3];
   }
-  return t0;
+  t0 = undefined;
+  let t1;
+  if ($[4] !== baseVideos || $[5] !== thumbnails) {
+    t1 = <FlatList baseVideos={baseVideos} items={thumbnails} />;
+    $[4] = baseVideos;
+    $[5] = thumbnails;
+    $[6] = t1;
+  } else {
+    t1 = $[6];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.expect.md
index 9bcc56f32f..9c87512a0f 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.expect.md
@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @validateNoSetStateInRender:false
 import {useMemo} from 'react';
 import {makeArray} from 'shared-runtime';
 
@@ -20,7 +21,7 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @validateNoSetStateInRender:false
 import { useMemo } from "react";
 import { makeArray } from "shared-runtime";
 
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.ts b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.ts
index cd84e42862..317491efbf 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.ts
@@ -1,3 +1,4 @@
+// @validateNoSetStateInRender:false
 import {useMemo} from 'react';
 import {makeArray} from 'shared-runtime';
 
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-nested-ifs.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-nested-ifs.expect.md
index f7b3605f4d..12f51643dd 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-nested-ifs.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-nested-ifs.expect.md
@@ -24,10 +24,12 @@ export const FIXTURE_ENTRYPOINT = {
 
 ```javascript
 function Component(props) {
+  let t0;
   if (props.cond) {
     if (props.cond) {
     }
   }
+  t0 = undefined;
 }
 
 export const FIXTURE_ENTRYPOINT = {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-return-empty.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-return-empty.expect.md
index be20ee39bd..b348ae34b6 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-return-empty.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-return-empty.expect.md
@@ -15,7 +15,10 @@ function component(a) {
 
 ```javascript
 function component(a) {
+  let t0;
+
   mutate(a);
+  t0 = undefined;
 }
 
 ```